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", ]