fix(lisong): fix icm/rnd+onppo config bugs and app_key env bugs

ding/entry/serial_entry_reward_model_offpolicy.py

@@ -14,6 +14,21 @@
 from ding.reward_model import create_reward_model
 from ding.utils import set_pkg_seed
 from .utils import random_collect
+import numpy as np
+from ding.utils import save_file
+
+
+def save_reward_model(path, reward_model, weights_name='best'):
+    path = os.path.join(path, 'reward_model', 'ckpt')
+    if not os.path.exists(path):
+        try:
+            os.makedirs(path)
+        except FileExistsError:
+            pass
+    path = os.path.join(path, 'ckpt_{}.pth.tar'.format(weights_name))
+    state_dict = reward_model.reward_model.state_dict()
+    save_file(path, state_dict)
+    print('Saved reward model ckpt in {}'.format(path))
 
 
 def serial_pipeline_reward_model_offpolicy(
@@ -87,11 +102,17 @@ def serial_pipeline_reward_model_offpolicy(
     # Accumulate plenty of data at the beginning of training.
     if cfg.policy.get('random_collect_size', 0) > 0:
         random_collect(cfg.policy, policy, collector, collector_env, commander, replay_buffer)
+    count = 0
+    best_reward = -np.inf
     while True:
         collect_kwargs = commander.step()
         # Evaluate policy performance
         if evaluator.should_eval(learner.train_iter):
             stop, reward = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep)
+            reward_mean = np.array([r['eval_episode_return'] for r in reward]).mean()
+            if reward_mean >= best_reward:
+                save_reward_model(cfg.exp_name, reward_model, 'best')
+                best_reward = reward_mean
             if stop:
                 break
         new_data_count, target_new_data_count = 0, cfg.reward_model.get('target_new_data_count', 1)
@@ -103,7 +124,9 @@ def serial_pipeline_reward_model_offpolicy(
             replay_buffer.push(new_data, cur_collector_envstep=collector.envstep)
         # update reward_model
         reward_model.train()
-        reward_model.clear_data()
+        # clear buffer per fix iters to make sure replay buffer's data count isn't too few.
+        if count % cfg.reward_model.clear_buffer_per_iters == 0:
+            reward_model.clear_data()
         # Learn policy from collected data
         for i in range(cfg.policy.learn.update_per_collect):
             # Learner will train ``update_per_collect`` times in one iteration.
@@ -122,7 +145,9 @@ def serial_pipeline_reward_model_offpolicy(
                 replay_buffer.update(learner.priority_info)
         if collector.envstep >= max_env_step or learner.train_iter >= max_train_iter:
             break
+        count += 1
 
     # Learner's after_run hook.
     learner.call_hook('after_run')
+    save_reward_model(cfg.exp_name, reward_model, 'last')
     return policy

ding/entry/serial_entry_reward_model_onpolicy.py

@@ -14,6 +14,21 @@
 from ding.reward_model import create_reward_model
 from ding.utils import set_pkg_seed
 from .utils import random_collect
+from ding.utils import save_file
+import numpy as np
+
+
+def save_reward_model(path, reward_model, weights_name='best'):
+    path = os.path.join(path, 'reward_model', 'ckpt')
+    if not os.path.exists(path):
+        try:
+            os.makedirs(path)
+        except FileExistsError:
+            pass
+    path = os.path.join(path, 'ckpt_{}.pth.tar'.format(weights_name))
+    state_dict = reward_model.reward_model.state_dict()
+    save_file(path, state_dict)
+    print('Saved reward model ckpt in {}'.format(path))
 
 
 def serial_pipeline_reward_model_onpolicy(
@@ -88,11 +103,16 @@ def serial_pipeline_reward_model_onpolicy(
     if cfg.policy.get('random_collect_size', 0) > 0:
         random_collect(cfg.policy, policy, collector, collector_env, commander, replay_buffer)
     count = 0
+    best_reward = -np.inf
     while True:
         collect_kwargs = commander.step()
         # Evaluate policy performance
         if evaluator.should_eval(learner.train_iter):
             stop, reward = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep)
+            reward_mean = np.array([r['eval_episode_return'] for r in reward]).mean()
+            if reward_mean >= best_reward:
+                save_reward_model(cfg.exp_name, reward_model, 'best')
+                best_reward = reward_mean
             if stop:
                 break
         new_data_count, target_new_data_count = 0, cfg.reward_model.get('target_new_data_count', 1)
@@ -127,4 +147,5 @@ def serial_pipeline_reward_model_onpolicy(
 
     # Learner's after_run hook.
     learner.call_hook('after_run')
+    save_reward_model(cfg.exp_name, reward_model, 'last')
     return policy

ding/reward_model/icm_reward_model.py

@@ -151,6 +151,9 @@ class ICMRewardModel(BaseRewardModel):
         update_per_collect=100,
         # (float) the importance weight of the forward and reverse loss
         reverse_scale=1,
+        intrinsic_reward_weight=0.003,  # 1/300
+        extrinsic_reward_norm=True,
+        extrinsic_reward_norm_max=1,
     )
 
     def __init__(self, config: EasyDict, device: str, tb_logger: 'SummaryWriter') -> None:  # noqa
@@ -171,8 +174,12 @@ def __init__(self, config: EasyDict, device: str, tb_logger: 'SummaryWriter') ->
         self.ce = nn.CrossEntropyLoss(reduction="mean")
         self.forward_mse = nn.MSELoss(reduction='none')
         self.reverse_scale = config.reverse_scale
+        self.res = nn.Softmax(dim=-1)
+        self.estimate_cnt_icm = 0
+        self.train_cnt_icm = 0
 
     def _train(self) -> None:
+        self.train_cnt_icm += 1
         train_data_list = [i for i in range(0, len(self.train_states))]
         train_data_index = random.sample(train_data_list, self.cfg.batch_size)
         data_states: list = [self.train_states[i] for i in train_data_index]
@@ -187,6 +194,13 @@ def _train(self) -> None:
         )
         inverse_loss = self.ce(pred_action_logit, data_actions.long())
         forward_loss = self.forward_mse(pred_next_state_feature, real_next_state_feature.detach()).mean()
+        self.tb_logger.add_scalar('icm_reward/forward_loss', forward_loss, self.train_cnt_icm)
+        self.tb_logger.add_scalar('icm_reward/inverse_loss', inverse_loss, self.train_cnt_icm)
+        action = torch.argmax(self.res(pred_action_logit), -1)
+        accuracy = torch.sum(action == data_actions.squeeze(-1)).item() / data_actions.shape[0]
+        self.tb_logger.add_scalar('icm_reward/action_accuracy', accuracy, self.train_cnt_icm)
+        loss = self.reverse_scale * inverse_loss + forward_loss
+        self.tb_logger.add_scalar('icm_reward/total_loss', loss, self.train_cnt_icm)
         loss = self.reverse_scale * inverse_loss + forward_loss
         self.opt.zero_grad()
         loss.backward()
@@ -195,7 +209,6 @@ def _train(self) -> None:
     def train(self) -> None:
         for _ in range(self.cfg.update_per_collect):
             self._train()
-        self.clear_data()
 
     def estimate(self, data: list) -> List[Dict]:
         # NOTE: deepcopy reward part of data is very important,
@@ -207,17 +220,32 @@ def estimate(self, data: list) -> List[Dict]:
         actions = torch.cat(actions).to(self.device)
         with torch.no_grad():
             real_next_state_feature, pred_next_state_feature, _ = self.reward_model(states, next_states, actions)
-            reward = self.forward_mse(real_next_state_feature, pred_next_state_feature).mean(dim=1)
-            reward = (reward - reward.min()) / (reward.max() - reward.min() + 1e-8)
-            reward = reward.to(train_data_augmented[0]['reward'].device)
-            reward = torch.chunk(reward, reward.shape[0], dim=0)
-        for item, rew in zip(train_data_augmented, reward):
+            raw_icm_reward = self.forward_mse(real_next_state_feature, pred_next_state_feature).mean(dim=1)
+            self.estimate_cnt_icm += 1
+            self.tb_logger.add_scalar('icm_reward/raw_icm_reward_max', raw_icm_reward.max(), self.estimate_cnt_icm)
+            self.tb_logger.add_scalar('icm_reward/raw_icm_reward_mean', raw_icm_reward.mean(), self.estimate_cnt_icm)
+            self.tb_logger.add_scalar('icm_reward/raw_icm_reward_min', raw_icm_reward.min(), self.estimate_cnt_icm)
+            self.tb_logger.add_scalar('icm_reward/raw_icm_reward_std', raw_icm_reward.std(), self.estimate_cnt_icm)
+            icm_reward = (raw_icm_reward - raw_icm_reward.min()) / (raw_icm_reward.max() - raw_icm_reward.min() + 1e-8)
+            self.tb_logger.add_scalar('icm_reward/icm_reward_max', icm_reward.max(), self.estimate_cnt_icm)
+            self.tb_logger.add_scalar('icm_reward/icm_reward_mean', icm_reward.mean(), self.estimate_cnt_icm)
+            self.tb_logger.add_scalar('icm_reward/icm_reward_min', icm_reward.min(), self.estimate_cnt_icm)
+            self.tb_logger.add_scalar('icm_reward/icm_reward_std', icm_reward.std(), self.estimate_cnt_icm)
+            icm_reward = (raw_icm_reward - raw_icm_reward.min()) / (raw_icm_reward.max() - raw_icm_reward.min() + 1e-8)
+            icm_reward = icm_reward.to(self.device)
+        for item, icm_rew in zip(train_data_augmented, icm_reward):
             if self.intrinsic_reward_type == 'add':
-                item['reward'] += rew
+                if self.cfg.extrinsic_reward_norm:
+                    item['reward'] = item[
+                        'reward'] / self.cfg.extrinsic_reward_norm_max + icm_rew * self.cfg.intrinsic_reward_weight
+                else:
+                    item['reward'] = item['reward'] + icm_rew * self.cfg.intrinsic_reward_weight
             elif self.intrinsic_reward_type == 'new':
-                item['intrinsic_reward'] = rew
+                item['intrinsic_reward'] = icm_rew
+                if self.cfg.extrinsic_reward_norm:
+                    item['reward'] = item['reward'] / self.cfg.extrinsic_reward_norm_max
             elif self.intrinsic_reward_type == 'assign':
-                item['reward'] = rew
+                item['reward'] = icm_rew
 
         return train_data_augmented
 

ding/reward_model/rnd_reward_model.py

@@ -12,7 +12,7 @@
 from .base_reward_model import BaseRewardModel
 from ding.utils import RunningMeanStd
 from ding.torch_utils.data_helper import to_tensor
-import copy
+import numpy as np
 
 
 def collect_states(iterator):
@@ -60,19 +60,15 @@ class RndRewardModel(BaseRewardModel):
         obs_norm=True,
         obs_norm_clamp_min=-1,
         obs_norm_clamp_max=1,
-        intrinsic_reward_weight=None,
         # means the relative weight of RND intrinsic_reward.
-        # If intrinsic_reward_weight=None, we will automatically set it based on
-        # the absolute value of the difference between max and min extrinsic reward in the sampled mini-batch
-        # please refer to  estimate() method for details.
-        intrinsic_reward_rescale=0.01,
-        # means the rescale value of RND intrinsic_reward only used when intrinsic_reward_weight is None
+        intrinsic_reward_weight=0.01,
+        extrinsic_reward_norm=True,
+        extrinsic_reward_norm_max=1,
     )
 
     def __init__(self, config: EasyDict, device: str = 'cpu', tb_logger: 'SummaryWriter' = None) -> None:  # noqa
         super(RndRewardModel, self).__init__()
         self.cfg = config
-        self.intrinsic_reward_rescale = self.cfg.intrinsic_reward_rescale
         assert device == "cpu" or device.startswith("cuda")
         self.device = device
         if tb_logger is None:  # TODO
@@ -87,6 +83,7 @@ def __init__(self, config: EasyDict, device: str = 'cpu', tb_logger: 'SummaryWri
         self.opt = optim.Adam(self.reward_model.predictor.parameters(), config.learning_rate)
         self._running_mean_std_rnd_reward = RunningMeanStd(epsilon=1e-4)
         self.estimate_cnt_rnd = 0
+        self.train_cnt_icm = 0
         self._running_mean_std_rnd_obs = RunningMeanStd(epsilon=1e-4)
 
     def _train(self) -> None:
@@ -102,13 +99,15 @@ def _train(self) -> None:
 
         predict_feature, target_feature = self.reward_model(train_data)
         loss = F.mse_loss(predict_feature, target_feature.detach())
+        self.tb_logger.add_scalar('rnd_reward/loss', loss, self.train_cnt_icm)
         self.opt.zero_grad()
         loss.backward()
         self.opt.step()
 
     def train(self) -> None:
         for _ in range(self.cfg.update_per_collect):
             self._train()
+            self.train_cnt_icm += 1
 
     def estimate(self, data: list) -> List[Dict]:
         """
@@ -132,14 +131,16 @@ def estimate(self, data: list) -> List[Dict]:
             self._running_mean_std_rnd_reward.update(mse.cpu().numpy())
 
             # Note: according to the min-max normalization, transform rnd reward to [0,1]
-            rnd_reward = (mse - mse.min()) / (mse.max() - mse.min() + 1e-11)
+            rnd_reward = (mse - mse.min()) / (mse.max() - mse.min() + 1e-8)
 
+            # save the rnd_reward statistics into tb_logger
             self.estimate_cnt_rnd += 1
             self.tb_logger.add_scalar('rnd_reward/rnd_reward_max', rnd_reward.max(), self.estimate_cnt_rnd)
             self.tb_logger.add_scalar('rnd_reward/rnd_reward_mean', rnd_reward.mean(), self.estimate_cnt_rnd)
             self.tb_logger.add_scalar('rnd_reward/rnd_reward_min', rnd_reward.min(), self.estimate_cnt_rnd)
+            self.tb_logger.add_scalar('rnd_reward/rnd_reward_std', rnd_reward.std(), self.estimate_cnt_rnd)
 
-            rnd_reward = rnd_reward.to(train_data_augmented[0]['reward'].device)
+            rnd_reward = rnd_reward.to(self.device)
             rnd_reward = torch.chunk(rnd_reward, rnd_reward.shape[0], dim=0)
         """
         NOTE: Following normalization approach to extrinsic reward seems be not reasonable,
@@ -148,30 +149,26 @@ def estimate(self, data: list) -> List[Dict]:
         # rewards = torch.stack([data[i]['reward'] for i in range(len(data))])
         # rewards = (rewards - torch.min(rewards)) / (torch.max(rewards) - torch.min(rewards))
 
-        # TODO(pu): how to set intrinsic_reward_rescale automatically?
-        if self.cfg.intrinsic_reward_weight is None:
-            """
-            NOTE: the following way of setting self.cfg.intrinsic_reward_weight is only suitable for the dense
-            reward env like lunarlander, not suitable for the dense reward env.
-            In sparse reward env, e.g. minigrid, if the agent reaches the goal, it obtain reward ~1, otherwise 0.
-            Thus, in sparse reward env, it's reasonable to set the intrinsic_reward_weight approximately equal to
-             the inverse of max_episode_steps.
-            """
-            self.cfg.intrinsic_reward_weight = self.intrinsic_reward_rescale * max(
-                1,
-                abs(
-                    max([train_data_augmented[i]['reward'] for i in range(len(train_data_augmented))]) -
-                    min([train_data_augmented[i]['reward'] for i in range(len(train_data_augmented))])
-                )
-            )
         for item, rnd_rew in zip(train_data_augmented, rnd_reward):
             if self.intrinsic_reward_type == 'add':
-                item['reward'] = item['reward'] + rnd_rew * self.cfg.intrinsic_reward_weight
+                if self.cfg.extrinsic_reward_norm:
+                    item['reward'] = item[
+                        'reward'] / self.cfg.extrinsic_reward_norm_max + rnd_rew * self.cfg.intrinsic_reward_weight
+                else:
+                    item['reward'] = item['reward'] + rnd_rew * self.cfg.intrinsic_reward_weight
             elif self.intrinsic_reward_type == 'new':
                 item['intrinsic_reward'] = rnd_rew
+                if self.cfg.extrinsic_reward_norm:
+                    item['reward'] = item['reward'] / self.cfg.extrinsic_reward_norm_max
             elif self.intrinsic_reward_type == 'assign':
                 item['reward'] = rnd_rew
 
+        # save the augmented_reward statistics into tb_logger
+        rew = [item['reward'].cpu().numpy() for item in train_data_augmented]
+        self.tb_logger.add_scalar('augmented_reward/reward_max', np.max(rew), self.estimate_cnt_rnd)
+        self.tb_logger.add_scalar('augmented_reward/reward_mean', np.mean(rew), self.estimate_cnt_rnd)
+        self.tb_logger.add_scalar('augmented_reward/reward_min', np.min(rew), self.estimate_cnt_rnd)
+        self.tb_logger.add_scalar('augmented_reward/reward_std', np.std(rew), self.estimate_cnt_rnd)
         return train_data_augmented
 
     def collect_data(self, data: list) -> None:

dizoo/minigrid/__init__.py

@@ -1,4 +1,4 @@
-from gym.envs.registration import register
+from gymnasium.envs.registration import register
 
 register(id='MiniGrid-AKTDT-7x7-1-v0', entry_point='dizoo.minigrid.envs:AppleKeyToDoorTreasure_7x7_1')
 
@@ -10,4 +10,6 @@
 
 register(id='MiniGrid-AKTDT-19x19-v0', entry_point='dizoo.minigrid.envs:AppleKeyToDoorTreasure_19x19')
 
-register(id='MiniGrid-AKTDT-19x19-3-v0', entry_point='dizoo.minigrid.envs:AppleKeyToDoorTreasure_19x19_3')
+register(id='MiniGrid-AKTDT-19x19-3-v0', entry_point='dizoo.minigrid.envs:AppleKeyToDoorTreasure_19x19_3')
+
+register(id='MiniGrid-NoisyTV-v0', entry_point='dizoo.minigrid.envs:NoisyTVEnv')