add drex to new entry

Author: ruoyuGao

ding/entry/tests/test_application_entry_drex_collect_data.py

@@ -70,5 +70,5 @@ def test_drex_collecting_data():
     args.cfg[0].bc_iteration = 1000  # for unittest
     args.cfg[1].policy.type = 'bc'
     drex_collecting_data(args=args)
-    os.popen('rm -rf {}'.format(expert_policy_state_dict_path))
-    os.popen('rm -rf {}'.format(args.cfg[0].reward_model.offline_data_path))
+    #os.popen('rm -rf {}'.format(expert_policy_state_dict_path))
+    #os.popen('rm -rf {}'.format(args.cfg[0].reward_model.offline_data_path))

ding/reward_model/drex_reward_model.py

@@ -1,6 +1,7 @@
 import copy
 from easydict import EasyDict
 import pickle
+import numpy as np
 
 from ding.utils import REWARD_MODEL_REGISTRY
 
@@ -77,11 +78,26 @@ def load_expert_data(self) -> None:
         """
         super(DrexRewardModel, self).load_expert_data()
 
-        with open(self.cfg.reward_model.offline_data_path + '/suboptimal_data.pkl', 'rb') as f:
+        with open(self.cfg.offline_data_path + '/suboptimal_data.pkl', 'rb') as f:
             self.demo_data = pickle.load(f)
 
     def train(self):
-        self._train()
+        # check if gpu available
+        device = self.device  # torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        # Assume that we are on a CUDA machine, then this should print a CUDA device:
+        self._logger.info("device: {}".format(device))
+        training_inputs, training_outputs = self.training_obs, self.training_labels
+
+        cum_loss = 0.0
+        training_data = list(zip(training_inputs, training_outputs))
+        for epoch in range(self.cfg.update_per_collect):
+            np.random.shuffle(training_data)
+            training_obs, training_labels = zip(*training_data)
+            cum_loss = self._train(training_obs, training_labels)
+            self.train_iter += 1
+            self._logger.info("[epoch {}] loss {}".format(epoch, cum_loss))
+            self.tb_logger.add_scalar("drex_reward/train_loss_iteration", cum_loss, self.train_iter)
+
         return_dict = self.pred_data(self.demo_data)
         res, pred_returns = return_dict['real'], return_dict['pred']
         self._logger.info("real: " + str(res))

ding/reward_model/ngu_reward_model.py

@@ -453,7 +453,7 @@ class NGURewardModel(BaseRewardModel):
     )
 
     def __init__(self, config: EasyDict, device: str, tb_logger: 'SummaryWriter') -> None:
-        super(NGURewardModel).__init__()
+        super(NGURewardModel, self).__init__()
         self.cfg = config
         self.tb_logger = tb_logger
         self.estimate_cnt = 0

dizoo/classic_control/cartpole/config/cartpole_drex_dqn_config.py

@@ -11,37 +11,43 @@
     ),
     reward_model=dict(
         type='drex',
+        exp_name='cartpole_drex_dqn_seed0',
         min_snippet_length=5,
         max_snippet_length=100,
         checkpoint_min=0,
-        checkpoint_max=1000,
-        checkpoint_step=1000,
+        checkpoint_max=760,
+        checkpoint_step=760,
         learning_rate=1e-5,
         update_per_collect=1,
         # path to expert models that generate demonstration data
         # Users should add their own model path here. Model path should lead to an exp_name.
         # Absolute path is recommended.
         # In DI-engine, it is ``exp_name``.
         # For example, if you want to use dqn to generate demos, you can use ``spaceinvaders_dqn``
-        expert_model_path='expert_model_path_placeholder',
+        expert_model_path='cartpole_dqn_seed0/ckpt/ckpt_best.pth.tar',
         # path to save reward model
         # Users should add their own model path here.
         # Absolute path is recommended.
         # For example, if you use ``spaceinvaders_drex``, then the reward model will be saved in this directory.
-        reward_model_path='reward_model_path_placeholder + ./spaceinvaders.params',
+        reward_model_path='cartpole_drex_dqn_seed0/cartpole.params',
         # path to save generated observations.
         # Users should add their own model path here.
         # Absolute path is recommended.
         # For example, if you use ``spaceinvaders_drex``, then all the generated data will be saved in this directory.
-        offline_data_path='offline_data_path_placeholder',
+        offline_data_path='cartpole_drex_dqn_seed0',
         # path to pretrained bc model. If omitted, bc will be trained instead.
         # Users should add their own model path here. Model path should lead to a model ckpt.
         # Absolute path is recommended.
-        bc_path='bc_path_placeholder',
+        # bc_path='bc_path_placeholder',
         # list of noises
         eps_list=[0, 0.5, 1],
         num_trajs_per_bin=20,
+        num_trajs=6,
+        num_snippets=6000,
         bc_iterations=6000,
+        hidden_size_list=[512, 64, 1],
+        obs_shape=4,
+        action_shape=2,
     ),
     policy=dict(
         cuda=False,
@@ -57,7 +63,13 @@
             batch_size=64,
             learning_rate=0.001,
         ),
-        collect=dict(n_sample=8, collector=dict(get_train_sample=False, )),
+        collect=dict(
+            n_sample=8,
+            collector=dict(
+                get_train_sample=False,
+                reward_shaping=False,
+            ),
+        ),
         eval=dict(evaluator=dict(eval_freq=40, )),
         other=dict(
             eps=dict(
@@ -66,7 +78,7 @@
                 end=0.1,
                 decay=10000,
             ),
-            replay_buffer=dict(replay_buffer_size=20000, ),
+            replay_buffer=dict(replay_buffer_size=200000, ),
         ),
     ),
 )
@@ -79,7 +91,24 @@
     ),
     env_manager=dict(type='subprocess'),
     policy=dict(type='dqn'),
-    collector=dict(type='episode'),
 )
 cartpole_drex_dqn_create_config = EasyDict(cartpole_drex_dqn_create_config)
 create_config = cartpole_drex_dqn_create_config
+
+if __name__ == "__main__":
+    import argparse
+    import torch
+    from ding.config import read_config
+    from ding.entry import drex_collecting_data
+    from ding.entry import serial_pipeline_reward_model_offpolicy
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--cfg', type=str, default='please enter abs path for this file')
+    parser.add_argument('--seed', type=int, default=0)
+    parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu')
+    args = parser.parse_args()
+    args.cfg = read_config(args.cfg)
+    args.cfg[1].policy.type = 'bc'
+    args.cfg[0].policy.collect.n_episode = 8
+    del args.cfg[0].policy.collect.n_sample
+    drex_collecting_data(args)
+    serial_pipeline_reward_model_offpolicy((main_config, create_config), pretrain_reward=True, cooptrain_reward=False)