Refactor env launcher scripts (#126)

Author: yuecideng

configs/gym/pour_water/gym_config_simple.json

@@ -3,6 +3,17 @@
     "max_episodes": 5,
     "env": {
         "events": {
+            "record_camera": {
+                "func": "record_camera_data",
+                "mode": "interval",
+                "interval_step": 1,
+                "params": {
+                    "name": "cam1",
+                    "resolution": [320, 240],
+                    "eye": [2, 0, 2],
+                    "target": [0.5, 0, 1]
+                }
+            },
             "random_light": {
                 "func": "randomize_light",
                 "mode": "interval",

embodichain/lab/gym/utils/gym_utils.py

@@ -17,6 +17,8 @@
 import numpy as np
 import torch
 import dexsim
+import argparse
+import gymnasium
 
 from typing import Dict, Any, List, Tuple, Union, Sequence
 from gymnasium import spaces
@@ -374,8 +376,8 @@ def config_to_cfg(config: dict, manager_modules: list = None) -> "EmbodiedEnvCfg
         ArticulationCfg,
         LightCfg,
     )
-    from embodichain.lab.gym.envs import EmbodiedEnvCfg
     from embodichain.lab.sim.sensors import SensorCfg
+    from embodichain.lab.gym.envs import EmbodiedEnvCfg
     from embodichain.lab.gym.envs.managers import (
         SceneEntityCfg,
         EventCfg,
@@ -701,3 +703,115 @@ def assign_data_to_dict(
 
     last_key = keys[-1]
     current_data[last_key] = value
+
+
+def add_env_launcher_args_to_parser(parser: argparse.ArgumentParser) -> None:
+    """Add common environment launcher arguments to an existing argparse parser.
+
+    This function adds the following arguments to the provided parser:
+        - --num_envs: Number of environments to run in parallel (default: 1)
+        - --device: Device to run the environment on (default: 'cpu')
+        - --headless: Whether to perform the simulation in headless mode (default: False)
+        - --enable_rt: Whether to use RTX rendering backend for the simulation (default: False)
+        - --gpu_id: The GPU ID to use for the simulation (default: 0)
+        - --filter_visual_rand: Whether to filter out visual randomization (default: False)
+        - --gym_config: Path to gym config file (default: '')
+        - --action_config: Path to action config file (default: None)
+        - --preview: Whether to preview the environment after launching (default: False)
+
+    Note:
+        1. In preview mode, the environment will be launched and keep running in a loop for user interaction.
+
+    Args:
+        parser (argparse.ArgumentParser): The parser to which arguments will be added.
+    """
+    parser.add_argument(
+        "--num_envs",
+        help="The number of environments to run in parallel.",
+        default=1,
+        type=int,
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cpu",
+        help="Device to run the environment on, e.g., 'cpu' or 'cuda'.",
+    )
+    parser.add_argument(
+        "--headless",
+        help="Whether to perform the simulation in headless mode.",
+        default=False,
+        action="store_true",
+    )
+    parser.add_argument(
+        "--enable_rt",
+        help="Whether to use RTX rendering backend for the simulation.",
+        default=False,
+        action="store_true",
+    )
+    parser.add_argument(
+        "--gpu_id",
+        help="The GPU ID to use for the simulation.",
+        default=0,
+        type=int,
+    )
+    parser.add_argument(
+        "--filter_visual_rand",
+        help="Whether to filter out visual randomization.",
+        default=False,
+        action="store_true",
+    )
+    parser.add_argument(
+        "--gym_config",
+        type=str,
+        help="Path to gym config file.",
+        default="",
+        required=True,
+    )
+    parser.add_argument(
+        "--action_config", type=str, help="Path to action config file.", default=None
+    )
+    parser.add_argument(
+        "--preview",
+        help="Whether to preview the environment after launching.",
+        default=False,
+        action="store_true",
+    )
+
+
+def build_env_cfg_from_args(
+    args: argparse.Namespace,
+) -> tuple["EmbodiedEnvCfg", dict, dict]:
+    """Build environment configuration from command-line arguments.
+
+    Args:
+        args (argparse.Namespace): The parsed command-line arguments.
+
+    Returns:
+        tuple[EmbodiedEnvCfg, dict, dict]: A tuple containing the environment configuration object,
+            the original gym configuration dictionary, and the action configuration dictionary.
+    """
+    from embodichain.utils.utility import load_json
+    from embodichain.lab.gym.envs import EmbodiedEnvCfg
+    from embodichain.lab.sim import SimulationManagerCfg
+
+    gym_config = load_json(args.gym_config)
+    cfg: EmbodiedEnvCfg = config_to_cfg(
+        gym_config, manager_modules=DEFAULT_MANAGER_MODULES
+    )
+    cfg.filter_visual_rand = args.filter_visual_rand
+
+    action_config = {}
+    if args.action_config is not None:
+        action_config = load_json(args.action_config)
+        action_config["action_config"] = action_config
+
+    cfg.num_envs = args.num_envs
+    cfg.sim_cfg = SimulationManagerCfg(
+        headless=args.headless,
+        sim_device=args.device,
+        enable_rt=args.enable_rt,
+        gpu_id=args.gpu_id,
+    )
+
+    return cfg, gym_config, action_config