update scripts for metadrive envs

liuzuxin · Aug 27, 2023 · f504dd4 · f504dd4
1 parent 6ede2c2
commit f504dd4
Show file tree

Hide file tree

Showing 12 changed files with 24 additions and 0 deletions.
diff --git a/examples/eval/eval_bc.py b/examples/eval/eval_bc.py
@@ -31,6 +31,8 @@ def eval(args: EvalConfig):
     if args.device == "cpu":
         torch.set_num_threads(args.threads)
 
+    if "Metadrive" in cfg["task"]:
+        import gym
     env = gym.make(cfg["task"])
     env.set_target_cost(cfg["cost_limit"])
 

diff --git a/examples/eval/eval_bcql.py b/examples/eval/eval_bcql.py
@@ -31,6 +31,8 @@ def eval(args: EvalConfig):
     if args.device == "cpu":
         torch.set_num_threads(args.threads)
 
+    if "Metadrive" in cfg["task"]:
+        import gym
     env = wrap_env(
         env=gym.make(cfg["task"]),
         reward_scale=cfg["reward_scale"],

diff --git a/examples/eval/eval_bearl.py b/examples/eval/eval_bearl.py
@@ -31,6 +31,8 @@ def eval(args: EvalConfig):
     if args.device == "cpu":
         torch.set_num_threads(args.threads)
 
+    if "Metadrive" in cfg["task"]:
+        import gym
     env = wrap_env(
         env=gym.make(cfg["task"]),
         reward_scale=cfg["reward_scale"],

diff --git a/examples/eval/eval_cdt.py b/examples/eval/eval_cdt.py
@@ -33,6 +33,8 @@ def eval(args: EvalConfig):
     if args.device == "cpu":
         torch.set_num_threads(args.threads)
 
+    if "Metadrive" in cfg["task"]:
+        import gym
     env = wrap_env(
         env=gym.make(cfg["task"]),
         reward_scale=cfg["reward_scale"],

diff --git a/examples/eval/eval_coptidice.py b/examples/eval/eval_coptidice.py
@@ -31,6 +31,8 @@ def eval(args: EvalConfig):
     if args.device == "cpu":
         torch.set_num_threads(args.threads)
 
+    if "Metadrive" in cfg["task"]:
+        import gym
     env = wrap_env(
         env=gym.make(cfg["task"]),
         reward_scale=cfg["reward_scale"],

diff --git a/examples/eval/eval_cpq.py b/examples/eval/eval_cpq.py
@@ -31,6 +31,8 @@ def eval(args: EvalConfig):
     if args.device == "cpu":
         torch.set_num_threads(args.threads)
 
+    if "Metadrive" in cfg["task"]:
+        import gym
     env = wrap_env(
         env=gym.make(cfg["task"]),
         reward_scale=cfg["reward_scale"],

diff --git a/examples/train/train_bc.py b/examples/train/train_bc.py
@@ -50,6 +50,8 @@ def train(args: BCTrainConfig):
         torch.set_num_threads(args.threads)
 
     # the cost scale is down in trainer rollout
+    if "Metadrive" in args.task:
+        import gym
     env = gym.make(args.task)
     data = env.get_dataset()
     env.set_target_cost(args.cost_limit)

diff --git a/examples/train/train_bcql.py b/examples/train/train_bcql.py
@@ -49,6 +49,8 @@ def train(args: BCQLTrainConfig):
         torch.set_num_threads(args.threads)
 
     # initialize environment
+    if "Metadrive" in args.task:
+        import gym
     env = gym.make(args.task)
 
     # pre-process offline dataset

diff --git a/examples/train/train_bearl.py b/examples/train/train_bearl.py
@@ -49,6 +49,8 @@ def train(args: BEARLTrainConfig):
         torch.set_num_threads(args.threads)
 
     # initialize environment
+    if "Metadrive" in args.task:
+        import gym
     env = gym.make(args.task)
 
     # pre-process offline dataset

diff --git a/examples/train/train_cdt.py b/examples/train/train_cdt.py
@@ -49,6 +49,8 @@ def train(args: CDTTrainConfig):
         torch.set_num_threads(args.threads)
 
     # initialize environment
+    if "Metadrive" in args.task:
+        import gym
     env = gym.make(args.task)
 
     # pre-process offline dataset

diff --git a/examples/train/train_coptidice.py b/examples/train/train_coptidice.py
@@ -50,6 +50,8 @@ def train(args: COptiDICETrainConfig):
         torch.set_num_threads(args.threads)
 
     # initialize environment
+    if "Metadrive" in args.task:
+        import gym
     env = gym.make(args.task)
 
     # pre-process offline dataset

diff --git a/examples/train/train_cpq.py b/examples/train/train_cpq.py
@@ -49,6 +49,8 @@ def train(args: CPQTrainConfig):
         torch.set_num_threads(args.threads)
 
     # initialize environment
+    if "Metadrive" in args.task:
+        import gym
     env = gym.make(args.task)
 
     # pre-process offline dataset