🚀 [RofuncRL] Update RofuncDTrans

Skylark0924 · Aug 30, 2023 · 4ef29f7 · 4ef29f7
1 parent 7c11d6c
commit 4ef29f7
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/examples/learning_rl/example_D4RL_RofuncRL.py b/examples/learning_rl/example_D4RL_RofuncRL.py
@@ -82,11 +82,11 @@ def inference(custom_args):
 
 
 if __name__ == '__main__':
-    gpu_id = 0
+    gpu_id = 3
 
     parser = argparse.ArgumentParser()
     # Available tasks: Hopper, HalfCheetah, Walker2d, Reacher2d
-    parser.add_argument("--task", type=str, default="Walker2d")
+    parser.add_argument("--task", type=str, default="Hopper")
     parser.add_argument("--agent", type=str, default="dtrans")  # dtrans
     parser.add_argument("--sim_device", type=str, default="cuda:{}".format(gpu_id))
     parser.add_argument("--rl_device", type=str, default="cuda:{}".format(gpu_id))

diff --git a/rofunc/learning/utils/download_datasets.py b/rofunc/learning/utils/download_datasets.py
@@ -17,7 +17,6 @@
 import pickle
 
 import gym
-import d4rl  # Import required to register environments, you may need to also import the submodule
 import numpy as np
 
 import rofunc as rf
@@ -34,6 +33,8 @@ def download_d4rl_dataset(save_dir):
             if os.path.exists(f'{save_dir}/{name}.pkl'):
                 continue
 
+            import d4rl  # Import required to register environments, you may need to also import the submodule
+
             env = gym.make(name)
             dataset = env.get_dataset()