wip

Author: yuecideng

docs/source/resources/robot/cobotmagic.md

@@ -39,9 +39,8 @@ CobotMagic is a versatile dual-arm collaborative robot developed by AgileX Robot
 from embodichain.lab.sim import SimulationManager, SimulationManagerCfg
 from embodichain.lab.sim.robots import CobotMagicCfg
 
-config = SimulationManagerCfg(headless=False, sim_device="cpu")
+config = SimulationManagerCfg(headless=False, sim_device="cpu", num_envs=2)
 sim = SimulationManager(config)
-sim.build_multiple_arenas(2)  # Supports parallel simulation in multiple arenas
 sim.set_manual_update(False)
 
 robot = sim.add_robot(cfg=CobotMagicCfg().from_dict({}))

docs/source/tutorial/create_scene.rst

@@ -31,15 +31,13 @@ Command-line arguments are parsed using ``argparse`` to allow for easy customiza
 .. literalinclude:: ../../../scripts/tutorials/sim/create_scene.py
    :language: python
    :start-at: # Parse command line arguments
-   :end-at: sim.build_multiple_arenas(args.num_envs, space=3.0)
+   :end-at: sim = SimulationManager(sim_cfg)
 
 There are two kinds of physics mode in :class:`SimulationManager`:
 
 - `manual`: The physics updates only when the user calls the :meth:`SimulationManager.update` function. This mode is used for robot learning tasks where precise control over simulation steps is required. Enabled by setting :meth:`SimulationManager.set_manual_update` to True.
 - `auto`: The physics updates in a standalone thread, which enable asynchronous rendering and physics stepping. This mode is suitable for visualizations and demos for digital twins applications. This is the default mode.
 
-If `num_envs` is greater than 1, :meth:`SimulationManager.build_multiple_arenas` should be used to create multiple simulation arenas.
-
 Adding objects to the scene
 ---------------------------
 

docs/source/tutorial/create_softbody.rst

@@ -32,8 +32,6 @@ The first step is to configure the simulation environment. This is done using th
    :start-at: # Configure the simulation
    :end-at:     print("[INFO]: Scene setup complete!")
 
-If ``num_envs`` is greater than 1, :meth:`SimulationManager.build_multiple_arenas` should be used to create multiple simulation arenas.
-
 Adding a soft body to the scene
 -------------------------------
 

docs/source/tutorial/rl.rst

@@ -182,7 +182,7 @@ To start training, run:
 
 .. code-block:: bash
 
-   python embodichain/agents/rl/train.py --config configs/agents/rl/push_cube/train_config.json
+   python -m embodichain.agents.rl.train --config configs/agents/rl/push_cube/train_config.json
 
 Outputs
 -------

embodichain/lab/gym/envs/base_env.py

@@ -208,8 +208,6 @@ def _setup_scene(self, **kwargs):
         logger.log_info(
             f"Initializing {self.num_envs} environments on {self.sim_cfg.sim_device}."
         )
-        if self.num_envs > 1:
-            self.sim.build_multiple_arenas(self.num_envs)
 
         self.robot = self._setup_robot(**kwargs)
         if self.robot is None:

embodichain/lab/sim/robots/cobotmagic.py

@@ -207,9 +207,8 @@ def build_pk_serial_chain(
 
     torch.set_printoptions(precision=5, sci_mode=False)
 
-    config = SimulationManagerCfg(headless=False, sim_device="cuda")
+    config = SimulationManagerCfg(headless=False, sim_device="cuda", num_envs=2)
     sim = SimulationManager(config)
-    sim.build_multiple_arenas(2)
 
     config = {
         "init_pos": [0.0, 0.0, 1.0],

embodichain/lab/sim/robots/dexforce_w1/cfg.py

@@ -353,7 +353,6 @@ def build_pk_serial_chain(
 
     config = SimulationManagerCfg(headless=True, sim_device="cpu")
     sim = SimulationManager(config)
-    sim.build_multiple_arenas(1)
 
     cfg = DexforceW1Cfg.from_dict(
         {

embodichain/lab/sim/sim_manager.py

@@ -125,6 +125,9 @@ class SimulationManagerCfg:
     cpu_num: int = 1
     """The number of CPU threads to use for the simulation engine."""
 
+    num_envs: int = 1
+    """The number of parallel environments (arenas) to simulate."""
+
     arena_space: float = 5.0
     """The distance between each arena when building multiple arenas."""
 
@@ -243,6 +246,8 @@ def __init__(
         # Set physics to manual update mode by default.
         self.set_manual_update(True)
 
+        self._build_multiple_arenas(sim_config.num_envs)
+
     @property
     def num_envs(self) -> int:
         """Get the number of arenas in the simulation.
@@ -500,7 +505,7 @@ def close_window(self) -> None:
         self._world.close_window()
         self.is_window_opened = False
 
-    def build_multiple_arenas(self, num: int, space: float | None = None) -> None:
+    def _build_multiple_arenas(self, num: int, space: float | None = None) -> None:
         """Build multiple arenas in a grid pattern.
 
         This interface is used for vectorized simulation.

examples/sim/demo/grasp_cup_to_caffe.py

@@ -83,11 +83,11 @@ def initialize_simulation(args) -> SimulationManager:
         sim_device=args.device,
         enable_rt=args.enable_rt,
         physics_dt=1.0 / 100.0,
+        num_envs=args.num_envs,
+        arena_space=2.5,
     )
     sim = SimulationManager(config)
 
-    sim.build_multiple_arenas(args.num_envs, space=2.5)
-
     if args.enable_rt:
         light = sim.add_light(
             cfg=LightCfg(

examples/sim/gizmo/gizmo_camera.py

@@ -40,9 +40,6 @@ def main():
     parser = argparse.ArgumentParser(
         description="Create and simulate a camera with gizmo in SimulationManager"
     )
-    parser.add_argument(
-        "--num_envs", type=int, default=1, help="Number of environments to simulate"
-    )
     parser.add_argument(
         "--device",
         type=str,
@@ -72,10 +69,6 @@ def main():
     sim = SimulationManager(sim_cfg)
     sim.set_manual_update(False)
 
-    # Build multiple arenas if requested
-    if args.num_envs > 1:
-        sim.build_multiple_arenas(args.num_envs, space=3.0)
-
     # Add some objects to the scene for camera to observe
     for i in range(5):
         cube_cfg = RigidObjectCfg(