refactor(gry): refactor reward model

ding/entry/serial_entry_reward_model_offpolicy.py

@@ -18,12 +18,14 @@
 
 
 def serial_pipeline_reward_model_offpolicy(
-        input_cfg: Union[str, Tuple[dict, dict]],
-        seed: int = 0,
-        env_setting: Optional[List[Any]] = None,
-        model: Optional[torch.nn.Module] = None,
-        max_train_iter: Optional[int] = int(1e10),
-        max_env_step: Optional[int] = int(1e10),
+    input_cfg: Union[str, Tuple[dict, dict]],
+    seed: int = 0,
+    env_setting: Optional[List[Any]] = None,
+    model: Optional[torch.nn.Module] = None,
+    max_train_iter: Optional[int] = int(1e10),
+    max_env_step: Optional[int] = int(1e10),
+    cooptrain_reward: Optional[bool] = True,
+    pretrain_reward: Optional[bool] = False,
 ) -> 'Policy':  # noqa
     """
     Overview:
@@ -78,6 +80,8 @@ def serial_pipeline_reward_model_offpolicy(
         cfg.policy.other.commander, learner, collector, evaluator, replay_buffer, policy.command_mode
     )
     reward_model = create_reward_model(cfg.reward_model, policy.collect_mode.get_attribute('device'), tb_logger)
+    if pretrain_reward:
+        reward_model.train()
 
     # ==========
     # Main loop
@@ -108,10 +112,11 @@ def serial_pipeline_reward_model_offpolicy(
             # collect data for reward_model training
             reward_model.collect_data(new_data)
             replay_buffer.push(new_data, cur_collector_envstep=collector.envstep)
-        # update reward_model
-        reward_model.train()
+        # 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 count % cfg.reward_model.clear_buffer_per_iters == 0:
+        if hasattr(cfg.reward_model, 'clear_buffer_per_iters') and 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):

ding/entry/serial_entry_reward_model_onpolicy.py

@@ -18,12 +18,14 @@
 
 
 def serial_pipeline_reward_model_onpolicy(
-        input_cfg: Union[str, Tuple[dict, dict]],
-        seed: int = 0,
-        env_setting: Optional[List[Any]] = None,
-        model: Optional[torch.nn.Module] = None,
-        max_train_iter: Optional[int] = int(1e10),
-        max_env_step: Optional[int] = int(1e10),
+    input_cfg: Union[str, Tuple[dict, dict]],
+    seed: int = 0,
+    env_setting: Optional[List[Any]] = None,
+    model: Optional[torch.nn.Module] = None,
+    max_train_iter: Optional[int] = int(1e10),
+    max_env_step: Optional[int] = int(1e10),
+    cooptrain_reward: Optional[bool] = True,
+    pretrain_reward: Optional[bool] = False,
 ) -> 'Policy':  # noqa
     """
     Overview:
@@ -78,7 +80,8 @@ def serial_pipeline_reward_model_onpolicy(
         cfg.policy.other.commander, learner, collector, evaluator, replay_buffer, policy.command_mode
     )
     reward_model = create_reward_model(cfg.reward_model, policy.collect_mode.get_attribute('device'), tb_logger)
-
+    if pretrain_reward:
+        reward_model.train()
     # ==========
     # Main loop
     # ==========
@@ -106,10 +109,12 @@ def serial_pipeline_reward_model_onpolicy(
             new_data = collector.collect(train_iter=learner.train_iter, policy_kwargs=collect_kwargs)
             new_data_count += len(new_data)
             # collect data for reward_model training
-            reward_model.collect_data(new_data)
+            if cooptrain_reward:
+                reward_model.collect_data(new_data)
         # update reward_model
-        reward_model.train()
-        if count % cfg.reward_model.clear_buffer_per_iters == 0:
+        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()
         # Learn policy from collected data
         for i in range(cfg.policy.learn.update_per_collect):

ding/entry/tests/test_serial_entry_reward_model.py

@@ -5,38 +5,54 @@
 from copy import deepcopy
 
 from dizoo.classic_control.cartpole.config.cartpole_dqn_config import cartpole_dqn_config, cartpole_dqn_create_config
+from dizoo.classic_control.cartpole.config.cartpole_trex_offppo_config import cartpole_trex_offppo_config,\
+     cartpole_trex_offppo_create_config
 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 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
 
 cfg = [
     {
         'type': 'pdeil',
         "alpha": 0.5,
         "discrete_action": False
-    },
-    {
+    }, {
         'type': 'gail',
         'input_size': 5,
-        'hidden_size': 64,
+        'hidden_size_list': [64],
         'batch_size': 64,
-    },
-    {
+    }, {
         'type': 'pwil',
         's_size': 4,
         'a_size': 2,
         'sample_size': 500,
-    },
-    {
+    }, {
         'type': 'red',
         'sample_size': 5000,
-        'input_size': 5,
-        'hidden_size': 64,
+        'obs_shape': 4,
+        'action_shape': 1,
+        'hidden_size_list': [64, 1],
         'update_per_collect': 200,
         'batch_size': 128,
-    },
+    }, {
+        'type': 'trex',
+        'exp_name': 'cartpole_trex_offppo_seed0',
+        'min_snippet_length': 5,
+        'max_snippet_length': 100,
+        'checkpoint_min': 0,
+        'checkpoint_max': 6,
+        'checkpoint_step': 6,
+        'learning_rate': 1e-5,
+        'update_per_collect': 1,
+        'expert_model_path': 'cartpole_ppo_offpolicy_seed0',
+        'data_path': 'abs data path',
+        'hidden_size_list': [512, 64, 1],
+        'obs_shape': 4,
+        'action_shape': 2,
+    }
 ]
 
 
@@ -51,9 +67,15 @@ def test_irl(reward_model_config):
     expert_data_path = 'expert_data.pkl'
     state_dict = expert_policy.collect_mode.state_dict()
     config = deepcopy(cartpole_ppo_offpolicy_config), deepcopy(cartpole_ppo_offpolicy_create_config)
-    collect_demo_data(
-        config, seed=0, state_dict=state_dict, expert_data_path=expert_data_path, collect_count=collect_count
-    )
+    if reward_model_config.type == 'trex':
+        trex_config = [deepcopy(cartpole_trex_offppo_config), deepcopy(cartpole_trex_offppo_create_config)]
+        trex_config[0].reward_model = reward_model_config
+        args = EasyDict({'cfg': deepcopy(trex_config), 'seed': 0, 'device': 'cpu'})
+        trex_collecting_data(args=args)
+    else:
+        collect_demo_data(
+            config, seed=0, state_dict=state_dict, expert_data_path=expert_data_path, collect_count=collect_count
+        )
     # irl + rl training
     cp_cartpole_dqn_config = deepcopy(cartpole_dqn_config)
     cp_cartpole_dqn_create_config = deepcopy(cartpole_dqn_create_config)
@@ -64,10 +86,18 @@ def test_irl(reward_model_config):
         reward_model_config['expert_data_path'] = expert_data_path
     cp_cartpole_dqn_config.reward_model = reward_model_config
     cp_cartpole_dqn_config.policy.collect.n_sample = 128
+    cooptrain_reward = True
+    pretrain_reward = False
+    if reward_model_config.type == 'trex':
+        cooptrain_reward = False
+        pretrain_reward = True
     serial_pipeline_reward_model_offpolicy(
-        (cp_cartpole_dqn_config, cp_cartpole_dqn_create_config), seed=0, max_train_iter=2
+        (cp_cartpole_dqn_config, cp_cartpole_dqn_create_config),
+        seed=0,
+        max_train_iter=2,
+        pretrain_reward=pretrain_reward,
+        cooptrain_reward=cooptrain_reward
     )
-
     os.popen("rm -rf ckpt_* log expert_data.pkl")
 
 

ding/reward_model/__init__.py

@@ -13,3 +13,5 @@
 from .guided_cost_reward_model import GuidedCostRewardModel
 from .ngu_reward_model import RndNGURewardModel, EpisodicNGURewardModel
 from .icm_reward_model import ICMRewardModel
+from .network import RepresentationNetwork, RNDNetwork, REDNetwork, GAILNetwork, ICMNetwork, GCLNetwork, TREXNetwork
+from .reword_model_utils import concat_state_action_pairs, combine_intrinsic_exterinsic_reward, obs_norm, collect_states