change clear data from pipeline to RM && add ngu to new entry

Author: ruoyuGao

ding/entry/serial_entry_reward_model_offpolicy.py

@@ -115,9 +115,8 @@ def serial_pipeline_reward_model_offpolicy(
         # update reward_model, when you want to train reward_model inloop
         if cooptrain_reward:
             reward_model.train()
-        # clear buffer per fix iters to make sure replay buffer's data count isn't too few.
-        if hasattr(cfg.reward_model, 'clear_buffer_per_iters') and count % cfg.reward_model.clear_buffer_per_iters == 0:
-            reward_model.clear_data()
+            # clear buffer per fix iters to make sure replay buffer's data count isn't too few.
+            reward_model.clear_data(iter=count)
         # 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.

ding/entry/serial_entry_reward_model_onpolicy.py

@@ -114,8 +114,7 @@ def serial_pipeline_reward_model_onpolicy(
         # update reward_model
         if cooptrain_reward:
             reward_model.train()
-        if hasattr(cfg.reward_model, 'clear_buffer_per_iters') and count % cfg.reward_model.clear_buffer_per_iters == 0:
-            reward_model.clear_data()
+            reward_model.clear_data(iter=count)
         # 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.

ding/entry/tests/test_serial_entry_reward_model.py

@@ -10,6 +10,7 @@
 from dizoo.classic_control.cartpole.config.cartpole_ppo_offpolicy_config import cartpole_ppo_offpolicy_config, cartpole_ppo_offpolicy_create_config  # noqa
 from dizoo.classic_control.cartpole.config.cartpole_rnd_onppo_config import cartpole_ppo_rnd_config, cartpole_ppo_rnd_create_config  # noqa
 from dizoo.classic_control.cartpole.config.cartpole_ppo_icm_config import cartpole_ppo_icm_config, cartpole_ppo_icm_create_config  # noqa
+from dizoo.classic_control.cartpole.config.cartpole_ngu_config import cartpole_ngu_config, cartpole_ngu_create_config
 from ding.entry import serial_pipeline, collect_demo_data, serial_pipeline_reward_model_offpolicy, \
     serial_pipeline_reward_model_onpolicy
 from ding.entry.application_entry_trex_collect_data import trex_collecting_data
@@ -116,3 +117,12 @@ def test_icm():
         serial_pipeline_reward_model_offpolicy(config, seed=0, max_train_iter=2)
     except Exception:
         assert False, "pipeline fail"
+
+
+@pytest.mark.unittest
+def test_ngu():
+    config = [deepcopy(cartpole_ngu_config), deepcopy(cartpole_ngu_create_config)]
+    try:
+        serial_pipeline_reward_model_offpolicy(config, seed=0, max_train_iter=2)
+    except Exception:
+        assert False, "pipeline fail"

ding/policy/ngu.py

@@ -431,7 +431,7 @@ def _init_collect(self) -> None:
         # epsilon=0.4, alpha=9
         self.eps = {i: 0.4 ** (1 + 8 * i / (self._cfg.collect.env_num - 1)) for i in range(self._cfg.collect.env_num)}
 
-    def _forward_collect(self, data: dict) -> dict:
+    def _forward_collect(self, data: dict, eps: Optional[float]) -> dict:
         r"""
         Overview:
             Collect output according to eps_greedy plugin

ding/reward_model/base_reward_model.py

@@ -60,7 +60,7 @@ def collect_data(self, data) -> None:
         raise NotImplementedError()
 
     @abstractmethod
-    def clear_data(self) -> None:
+    def clear_data(self, iter: int) -> None:
         """
         Overview:
             Clearing training data. \

ding/reward_model/gail_irl_model.py

@@ -212,10 +212,14 @@ def collect_data(self, data: list) -> None:
         data = torch.unbind(data, dim=0)
         self.train_data.extend(data)
 
-    def clear_data(self) -> None:
+    def clear_data(self, iter: int) -> None:
         """
         Overview:
             Clearing training data. \
             This is a side effect function which clears the data attribute in ``self``
         """
-        self.train_data.clear()
+        assert hasattr(
+            self.cfg, 'clear_buffer_per_iters'
+        ), "Reward Model does not have clear_buffer_per_iters, Clear failed"
+        if iter % self.cfg.clear_buffer_per_iters == 0:
+            self.train_data.clear()

ding/reward_model/guided_cost_reward_model.py

@@ -132,7 +132,7 @@ def collect_data(self, data) -> None:
         # if online_net is trained continuously, there should be some implementations in collect_data method
         pass
 
-    def clear_data(self):
+    def clear_data(self, iter: int):
         """
         Overview:
             Collecting clearing data, not implemented if reward model (i.e. online_net) is only trained ones, \

ding/reward_model/icm_reward_model.py

@@ -175,11 +175,15 @@ def collect_data(self, data: list) -> None:
         self.train_next_states.extend(next_states)
         self.train_actions.extend(actions)
 
-    def clear_data(self) -> None:
-        self.train_data.clear()
-        self.train_states.clear()
-        self.train_next_states.clear()
-        self.train_actions.clear()
+    def clear_data(self, iter: int) -> None:
+        assert hasattr(
+            self.cfg, 'clear_buffer_per_iters'
+        ), "Reward Model does not have clear_buffer_per_iters, Clear failed"
+        if iter % self.cfg.clear_buffer_per_iters == 0:
+            self.train_data.clear()
+            self.train_states.clear()
+            self.train_next_states.clear()
+            self.train_actions.clear()
 
     def state_dict(self) -> Dict:
         return self.reward_model.state_dict()

ding/reward_model/ngu_reward_model.py

@@ -1,5 +1,6 @@
 import copy
 import random
+from typing import Any
 
 import numpy as np
 import torch
@@ -408,3 +409,92 @@ def fusion_reward(
                                     int(data[i]['beta'][j])]
 
         return data, estimate_cnt
+
+
+@REWARD_MODEL_REGISTRY.register('ngu-reward')
+class NGURewardModel(BaseRewardModel):
+    r"""
+    Overview:
+        The unifying reward for ngu which combined rnd-ngu and episodic
+        The corresponding paper is `never give up: learning directed exploration strategies`.
+    """
+    config = dict(
+        type='ngu-reward',
+        policy_nstep=5,
+        collect_env_num=8,
+        rnd_reward_model=dict(
+            intrinsic_reward_type='add',
+            learning_rate=5e-4,
+            obs_shape=4,
+            action_shape=2,
+            batch_size=128,  # transitions
+            update_per_collect=10,
+            only_use_last_five_frames_for_icm_rnd=False,
+            clear_buffer_per_iters=10,
+            nstep=5,
+            hidden_size_list=[128, 128, 64],
+            type='rnd-ngu',
+        ),
+        episodic_reward_model=dict(
+            last_nonzero_reward_rescale=False,
+            last_nonzero_reward_weight=1,
+            intrinsic_reward_type='add',
+            learning_rate=5e-4,
+            obs_shape=4,
+            action_shape=2,
+            batch_size=128,  # transitions
+            update_per_collect=10,
+            only_use_last_five_frames_for_icm_rnd=False,
+            clear_buffer_per_iters=10,
+            nstep=5,
+            hidden_size_list=[128, 128, 64],
+            type='episodic',
+        ),
+    )
+
+    def __init__(self, config: EasyDict, device: str, tb_logger: 'SummaryWriter') -> None:
+        super(NGURewardModel).__init__()
+        self.cfg = config
+        self.tb_logger = tb_logger
+        self.estimate_cnt = 0
+        self.rnd_reward_model = RndNGURewardModel(config.rnd_reward_model, device, tb_logger)
+        self.episodic_reward_model = EpisodicNGURewardModel(config.episodic_reward_model, device, tb_logger)
+
+    def train(self) -> None:
+        self.rnd_reward_model.train()
+        self.episodic_reward_model.train()
+
+    def estimate(self, data: list) -> dict:
+
+        # estimate reward
+        rnd_reward = self.rnd_reward_model.estimate(data)
+        episodic_reward = self.episodic_reward_model.estimate(data)
+
+        # combine reward
+        train_data_augumented, self.estimate_cnt = self.episodic_reward_model.fusion_reward(
+            data,
+            episodic_reward,
+            rnd_reward,
+            nstep=self.cfg.policy_nstep,
+            collector_env_num=self.cfg.collect_env_num,
+            tb_logger=self.tb_logger,
+            estimate_cnt=self.estimate_cnt
+        )
+
+        return train_data_augumented
+
+    def collect_data(self, data) -> None:
+        self.rnd_reward_model.collect_data(data)
+        self.episodic_reward_model.collect_data(data)
+
+    def clear_data(self, iter: int) -> None:
+        assert hasattr(
+            self.cfg.rnd_reward_model, 'clear_buffer_per_iters'
+        ), "RND Reward Model does not have clear_buffer_per_iters, Clear failed"
+        assert hasattr(
+            self.cfg.episodic_reward_model, 'clear_buffer_per_iters'
+        ), "Episodic Reward Model does not have clear_buffer_per_iters, Clear failed"
+        if iter % self.cfg.rnd_reward_model.clear_buffer_per_iters == 0:
+            self.rnd_reward_model.clear_data()
+        if iter % self.cfg.episodic_reward_model.clear_buffer_per_iters == 0:
+            self.episodic_reward_model.clear_data()

ding/reward_model/pdeil_irl_model.py

@@ -210,10 +210,14 @@ def collect_data(self, item: list):
         """
         self.train_data.extend(item)
 
-    def clear_data(self):
+    def clear_data(self, iter: int):
         """
         Overview:
             Clearing training data. \
             This is a side effect function which clears the data attribute in ``self``
         """
-        self.train_data.clear()
+        assert hasattr(
+            self.cfg, 'clear_buffer_per_iters'
+        ), "Reward Model does not have clear_buffer_per_iters, Clear failed"
+        if iter % self.cfg.clear_buffer_per_iters == 0:
+            self.train_data.clear()