rtg datamodule fixed

Author: Yanchao Sun

README.md

@@ -170,8 +170,8 @@ python src/dmc_downstream.py \
         data.num_steps=1000000 \
         domain_and_task.domain=${DOMAIN} \
         domain_and_task.task=${TASK} \
-        model.model_type=reward_conditioned
-        data.rand_select=False
+        model.model_type=reward_conditioned \
+        data.rand_select=True \
         no_load_action=True \
         load_model_from=./outputs/pretrain_explore/checkpoints/last.ckpt \
         output_dir=./outputs/${DOMAIN}_${TASK}_bc/

configs/downstream.yaml

@@ -66,10 +66,8 @@ data:
   num_workers: 1
   train_replay_id: 1
   val_replay_id: 2
-
-  # used for DMCBCDataModule only (used when model_type is "naive")
-  select_rate: 0
-  rand_select: False
+  select_rate: 0.1
+  rand_select: False # True for reward-conditioned learning
 
 # DTModel
 model:

src/datamodules/dmc_datamodule.py

@@ -240,20 +240,7 @@ def __getitem__(self, idx):
 
 
 class DMCDataModule(LightningDataModule):
-    """Example of LightningDataModule for Atari DQN Replay Buffer dataset.
-
-    A DataModule implements 5 key methods:
-        - prepare_data (things to do on 1 GPU/TPU, not on every GPU/TPU in distributed mode)
-        - setup (things to do on every accelerator in distributed mode)
-        - train_dataloader (the training dataloader)
-        - val_dataloader (the validation dataloader(s))
-        - test_dataloader (the test dataloader(s))
-
-    This allows you to share a full dataset without explaining how to download,
-    split, transform and process the data.
-
-    Read the docs:
-        https://pytorch-lightning.readthedocs.io/en/latest/extensions/datamodules.html
+    """DMC Dataset for one specific domain and task.
     """
 
     def __init__(
@@ -369,32 +356,12 @@ def val_dataloader(self):
 
     def test_dataloader(self):
         return None
-        # return DataLoader(
-        #     self.test_dataset,
-        #     shuffle=False,
-        #     pin_memory=True,
-        #     batch_size=self.hparams.batch_size,
-        #     num_workers=self.hparams.num_workers,
-        #     persistent_workers=True
-        # )
 
 
 class DMCBCDataModule(LightningDataModule):
-    """Example of LightningDataModule for behavior cloning on DMC. Only top 10 percent trajectories
-    are used.
-
-    A DataModule implements 5 key methods:
-        - prepare_data (things to do on 1 GPU/TPU, not on every GPU/TPU in distributed mode)
-        - setup (things to do on every accelerator in distributed mode)
-        - train_dataloader (the training dataloader)
-        - val_dataloader (the validation dataloader(s))
-        - test_dataloader (the test dataloader(s))
-
-    This allows you to share a full dataset without explaining how to download,
-    split, transform and process the data.
-
-    Read the docs:
-        https://pytorch-lightning.readthedocs.io/en/latest/extensions/datamodules.html
+    """DMC dataset on one specific domain and task with data subsampling. 
+    For behavior cloning: only top k percent trajectories (expert data) are used.
+    For reward-conditioned learning: randomly selected k percent trajectories are used.
     """
 
     def __init__(
@@ -412,7 +379,7 @@ def __init__(
         train_replay_id=1,
         val_replay_id=2,
         select_rate=0.1,
-        rand_select=False,
+        rand_select=False, # True for reward-conditioned learning
     ):
         super().__init__()
 

src/dmc_downstream.py

@@ -13,11 +13,7 @@
 def main(cfg):
     cfg_dict = OmegaConf.to_container(cfg, resolve=True)
 
-    if cfg.model.model_type == "naive":
-        # bc = True
-        dmc_data = DMCBCDataModule(**cfg.data)
-    else:
-        dmc_data = DMCDataModule(**cfg.data)
+    dmc_data = DMCBCDataModule(**cfg.data)
 
     model = pl_utils.instantiate_class(cfg["model"])
     if cfg.load_model_from: