git.diff

diff --git a/data/envs/metaworld/train.py b/data/envs/metaworld/train.py
index e01df51..f0d2bb9 100644
--- a/data/envs/metaworld/train.py
+++ b/data/envs/metaworld/train.py
@@ -1,30 +1,78 @@
-from typing import Optional
-import argparse
 import sys
-import metaworld
-import gym
+from typing import Optional
 
+import gym
+import metaworld
 from sample_factory.cfg.arguments import parse_full_cfg, parse_sf_args
 from sample_factory.envs.env_utils import register_env
 from sample_factory.train import run_rl
 
-
-def make_custom_env(full_env_name: str, cfg=None, env_config=None, render_mode: Optional[str] = None):
-    # see the section below explaining arguments
-    return gym.make("pick-place-v2", render_mode=render_mode)
+ENV_NAMES = [
+    "assembly-v2",
+    "basketball-v2",
+    "bin-picking-v2",
+    "box-close-v2",
+    "button-press-topdown-v2",
+    "button-press-topdown-wall-v2",
+    "button-press-v2",
+    "button-press-wall-v2",
+    "coffee-button-v2",
+    "coffee-pull-v2",
+    "coffee-push-v2",
+    "dial-turn-v2",
+    "disassemble-v2",
+    "door-close-v2",
+    "door-lock-v2",
+    "door-open-v2",
+    "door-unlock-v2",
+    "hand-insert-v2",
+    "drawer-close-v2",
+    "drawer-open-v2",
+    "faucet-open-v2",
+    "faucet-close-v2",
+    "hammer-v2",
+    "handle-press-side-v2",
+    "handle-press-v2",
+    "handle-pull-side-v2",
+    "handle-pull-v2",
+    "lever-pull-v2",
+    "peg-insert-side-v2",
+    "pick-place-wall-v2",
+    "pick-out-of-hole-v2",
+    "reach-v2",
+    "push-back-v2",
+    "push-v2",
+    "pick-place-v2",
+    "plate-slide-v2",
+    "plate-slide-side-v2",
+    "plate-slide-back-v2",
+    "plate-slide-back-side-v2",
+    "peg-insert-side-v2",
+    "peg-unplug-side-v2",
+    "soccer-v2",
+    "stick-push-v2",
+    "stick-pull-v2",
+    "push-wall-v2",
+    "push-v2",
+    "reach-wall-v2",
+    "reach-v2",
+    "shelf-place-v2",
+    "sweep-into-v2",
+    "sweep-v2",
+    "window-open-v2",
+    "window-close-v2",
+]
 
 
-def parse_args(argv=None, evaluation=False):
-    # parse the command line arguments to build
-    parser, partial_cfg = parse_sf_args(argv=argv, evaluation=evaluation)
-    final_cfg = parse_full_cfg(parser, argv)
-    return final_cfg
+def make_custom_env(full_env_name: str, cfg=None, env_config=None, render_mode: Optional[str] = None):
+    return gym.make(full_env_name, render_mode=render_mode)
 
 
 def main():
-    """Script entry point."""
-    register_env("pick-place-v2", make_custom_env)
-    cfg = parse_args()
+    for env_name in ENV_NAMES:
+        register_env(env_name, make_custom_env)
+    parser, _ = parse_sf_args(argv=None, evaluation=False)
+    cfg = parse_full_cfg(parser)
     status = run_rl(cfg)
     return status
 
diff --git a/setup.py b/setup.py
index 4cbcc4b..b35d687 100644
--- a/setup.py
+++ b/setup.py
@@ -65,7 +65,7 @@ REQUIRED_PKGS = [
     "huggingface_hub>=0.10",  # For sharing objects, environments & trained RL policies
     "gym==0.26.2",  # For RL action spaces and API
     "hydra-core",
-    "envpool",
+    # "envpool",
 ]