From e12a5940bffe39db40898dcc5f12febd3ae38290 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Mon, 9 Feb 2026 23:26:46 +0800 Subject: [PATCH 1/3] wip --- configs/gym/pour_water/gym_config_simple.json | 11 ++ embodichain/lab/gym/utils/gym_utils.py | 112 +++++++++++++- embodichain/lab/scripts/preview_env.py | 143 ------------------ embodichain/lab/scripts/run_agent.py | 79 +--------- embodichain/lab/scripts/run_env.py | 130 ++++++---------- 5 files changed, 175 insertions(+), 300 deletions(-) delete mode 100644 embodichain/lab/scripts/preview_env.py diff --git a/configs/gym/pour_water/gym_config_simple.json b/configs/gym/pour_water/gym_config_simple.json index ddde2e4d..3b709145 100644 --- a/configs/gym/pour_water/gym_config_simple.json +++ b/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", diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index 3555322c..ad3b8a5e 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/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 @@ -352,8 +354,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, @@ -679,3 +681,111 @@ 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="" + ) + 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 diff --git a/embodichain/lab/scripts/preview_env.py b/embodichain/lab/scripts/preview_env.py deleted file mode 100644 index 09f7cda3..00000000 --- a/embodichain/lab/scripts/preview_env.py +++ /dev/null @@ -1,143 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -import torch -import gymnasium -import argparse -import numpy as np - -from embodichain.lab.sim import SimulationManagerCfg -from embodichain.lab.gym.envs import EmbodiedEnvCfg -from embodichain.lab.gym.utils.gym_utils import ( - config_to_cfg, - DEFAULT_MANAGER_MODULES, -) -from embodichain.utils.utility import load_json -from embodichain.utils import logger - -if __name__ == "__main__": - np.set_printoptions(precision=5, suppress=True) - torch.set_printoptions(precision=5, sci_mode=False) - - parser = argparse.ArgumentParser() - 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("--gym_config", type=str, help="gym_config", default="") - parser.add_argument( - "--action_config", - type=str, - help="Path to the action configuration file.", - default=None, - ) - parser.add_argument( - "--filter_visual_rand", - help="Whether to filter out visual randomization.", - default=False, - action="store_true", - ) - - args = parser.parse_args() - - """ - TODO: Currently, this file is only used to preview the template.json config file. - We may add more features to support more general case parsing from config files. - """ - - ############################################################################################## - # load gym config - 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, - ) - - env = gymnasium.make(id=gym_config["id"], cfg=cfg, **action_config) - - obs, info = env.reset() - - """ - Run the following code to create a demonstration and perform env steps. - - ``` - # Demo version of environment rollout - for i in range(10): - qpos = env.robot.get_qpos() - - obs, reward, done, truncated, info = env.step(qpos) - - # reset the environment - env.reset() - ``` - - Run the following code to preview the sensor observations. - - ``` - env.preview_sensor_data("camera") - ``` - """ - - end = False - while end is False: - print("Press `p` to into embed mode to interact with the environment.") - print("Press `q` to quit the simulation.") - txt = input() - if txt == "p": - from IPython import embed - - embed() - elif txt == "q": - end = True diff --git a/embodichain/lab/scripts/run_agent.py b/embodichain/lab/scripts/run_agent.py index f8e248ec..c635d4a2 100644 --- a/embodichain/lab/scripts/run_agent.py +++ b/embodichain/lab/scripts/run_agent.py @@ -20,11 +20,9 @@ import torch from embodichain.utils.utility import load_json -from embodichain.lab.sim import SimulationManagerCfg -from embodichain.lab.gym.envs import EmbodiedEnvCfg from embodichain.lab.gym.utils.gym_utils import ( - config_to_cfg, - DEFAULT_MANAGER_MODULES, + add_env_launcher_args_to_parser, + build_env_cfg_from_args, ) from embodichain.utils.logger import log_warning, log_info, log_error from .run_env import main @@ -35,61 +33,7 @@ torch.set_printoptions(precision=5, sci_mode=False) parser = argparse.ArgumentParser() - 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( - "--render_backend", - help="The rendering backend to use for the simulation.", - default="egl", - type=str, - ) - parser.add_argument( - "--gpu_id", - help="The GPU ID to use for the simulation.", - default=0, - type=int, - ) - parser.add_argument( - "--debug_mode", - help="Enable debug mode.", - default=False, - action="store_true", - ) - 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 the gym configuration file.", - required=True, - ) + add_env_launcher_args_to_parser(parser) parser.add_argument( "--task_name", type=str, @@ -117,26 +61,13 @@ exit(1) # Load configurations - gym_config = load_json(args.gym_config) + env_cfg, gym_config, action_config = build_env_cfg_from_args(args) agent_config = load_json(args.agent_config) - # Build environment configuration - cfg: EmbodiedEnvCfg = config_to_cfg( - gym_config, manager_modules=DEFAULT_MANAGER_MODULES - ) - cfg.filter_visual_rand = args.filter_visual_rand - 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, - ) - # Create environment env = gymnasium.make( id=gym_config["id"], - cfg=cfg, + cfg=env_cfg, agent_config=agent_config, agent_config_path=args.agent_config, task_name=args.task_name, diff --git a/embodichain/lab/scripts/run_env.py b/embodichain/lab/scripts/run_env.py index 6b3844b1..c5b36df9 100644 --- a/embodichain/lab/scripts/run_env.py +++ b/embodichain/lab/scripts/run_env.py @@ -21,14 +21,9 @@ import torch import tqdm -from threading import Thread - -from embodichain.utils.utility import load_json -from embodichain.lab.sim import SimulationManagerCfg -from embodichain.lab.gym.envs import EmbodiedEnvCfg from embodichain.lab.gym.utils.gym_utils import ( - config_to_cfg, - DEFAULT_MANAGER_MODULES, + add_env_launcher_args_to_parser, + build_env_cfg_from_args, ) from embodichain.utils.logger import log_warning, log_info, log_error @@ -111,6 +106,9 @@ def generate_function( def main(args, env, gym_config): + if args.preview: + log_warning("Preview mode enabled. Launching environment preview...") + preview(env) log_info("Start offline data generation.", color="green") # TODO: Support multiple trajectories per episode generation. @@ -122,93 +120,61 @@ def main(args, env, gym_config): i, save_path=getattr(args, "save_path", ""), save_video=getattr(args, "save_video", False), - debug_mode=args.debug_mode, + debug_mode=getattr(args, "debug_mode", False), regenerate=getattr(args, "regenerate", False), ) +def preview(env: gymnasium.Env) -> None: + """ + Run the following code to create a demonstration and perform env steps. + + ``` + # Demo version of environment rollout + for i in range(10): + qpos = env.robot.get_qpos() + + obs, reward, done, truncated, info = env.step(qpos) + + # reset the environment + env.reset() + ``` + + Run the following code to preview the sensor observations. + + ``` + env.preview_sensor_data("camera") + ``` + """ + _, _ = env.reset() + + end = False + while end is False: + print("Press `p` to into embed mode to interact with the environment.") + print("Press `q` to quit the simulation.") + txt = input() + if txt == "p": + from IPython import embed + + embed() + elif txt == "q": + end = True + + exit(0) + + if __name__ == "__main__": np.set_printoptions(5, suppress=True) torch.set_printoptions(precision=5, sci_mode=False) parser = argparse.ArgumentParser() - 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( - "--save_path", help="path", default="./outputs/thirdviewvideo", type=str - ) - parser.add_argument( - "--save_video", - help="Whether to save the video of the episode.", - default=False, - action="store_true", - ) - parser.add_argument( - "--debug_mode", - help="Enable debug mode.", - default=False, - action="store_true", - ) - 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="gym_config", default="") - parser.add_argument("--action_config", type=str, help="action_config", default=None) + + add_env_launcher_args_to_parser(parser) args = parser.parse_args() - # if args.num_envs != 1: - # log_error(f"Currently only support num_envs=1, but got {args.num_envs}.") + env_cfg, gym_config, action_config = build_env_cfg_from_args(args) - 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, - ) + env = gymnasium.make(id=gym_config["id"], cfg=env_cfg, **action_config) - env = gymnasium.make(id=gym_config["id"], cfg=cfg, **action_config) main(args, env, gym_config) From 4b95a4798073afa11f60e8aa85de19b80712478e Mon Sep 17 00:00:00 2001 From: yuecideng Date: Mon, 9 Feb 2026 23:51:24 +0800 Subject: [PATCH 2/3] wip --- embodichain/lab/gym/utils/gym_utils.py | 2 +- embodichain/lab/scripts/run_env.py | 11 +++++++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index ad3b8a5e..b944fd54 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -740,7 +740,7 @@ def add_env_launcher_args_to_parser(parser: argparse.ArgumentParser) -> None: action="store_true", ) parser.add_argument( - "--gym_config", type=str, help="Path to gym config file.", default="" + "--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 diff --git a/embodichain/lab/scripts/run_env.py b/embodichain/lab/scripts/run_env.py index c5b36df9..1804be66 100644 --- a/embodichain/lab/scripts/run_env.py +++ b/embodichain/lab/scripts/run_env.py @@ -150,11 +150,18 @@ def preview(env: gymnasium.Env) -> None: end = False while end is False: - print("Press `p` to into embed mode to interact with the environment.") + print("Press `p` to enter embed mode to interact with the environment.") print("Press `q` to quit the simulation.") txt = input() if txt == "p": - from IPython import embed + try: + from IPython import embed + except ImportError: + log_error( + "IPython is not installed. Preview mode requires IPython to be " + "available. Please install it with `pip install ipython` and try again." + ) + continue embed() elif txt == "q": From fa9587764b08b4165997a8da8b4b6dd37e9bc0e8 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Mon, 9 Feb 2026 23:54:37 +0800 Subject: [PATCH 3/3] wip --- embodichain/lab/gym/utils/gym_utils.py | 6 +++++- embodichain/lab/scripts/run_agent.py | 2 +- embodichain/lab/scripts/run_env.py | 4 ++-- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index d3d23c39..2d2af21a 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -762,7 +762,11 @@ def add_env_launcher_args_to_parser(parser: argparse.ArgumentParser) -> None: action="store_true", ) parser.add_argument( - "--gym_config", type=str, help="Path to gym config file.", default="", required=True + "--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 diff --git a/embodichain/lab/scripts/run_agent.py b/embodichain/lab/scripts/run_agent.py index c635d4a2..54c14768 100644 --- a/embodichain/lab/scripts/run_agent.py +++ b/embodichain/lab/scripts/run_agent.py @@ -24,7 +24,7 @@ add_env_launcher_args_to_parser, build_env_cfg_from_args, ) -from embodichain.utils.logger import log_warning, log_info, log_error +from embodichain.utils.logger import log_error from .run_env import main diff --git a/embodichain/lab/scripts/run_env.py b/embodichain/lab/scripts/run_env.py index 1804be66..858950af 100644 --- a/embodichain/lab/scripts/run_env.py +++ b/embodichain/lab/scripts/run_env.py @@ -106,7 +106,7 @@ def generate_function( def main(args, env, gym_config): - if args.preview: + if getattr(args, "preview", False): log_warning("Preview mode enabled. Launching environment preview...") preview(env) @@ -134,7 +134,7 @@ def preview(env: gymnasium.Env) -> None: for i in range(10): qpos = env.robot.get_qpos() - obs, reward, done, truncated, info = env.step(qpos) + obs, reward, terminated, truncated, info = env.step(qpos) # reset the environment env.reset()