Creating a Basic Environment

Creating a Basic Environment#

This tutorial shows you how to create a simple robot learning environment using EmbodiChain’s Gym interface. You’ll learn how to inherit from the base environment class, set up robots and objects, define actions and observations, and run training scenarios.

The Code#

The tutorial corresponds to the random_reach.py script in the scripts/tutorials/gym directory.

Code for random_reach.py
  1# ----------------------------------------------------------------------------
  2# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
  3#
  4# Licensed under the Apache License, Version 2.0 (the "License");
  5# you may not use this file except in compliance with the License.
  6# You may obtain a copy of the License at
  7#
  8#     http://www.apache.org/licenses/LICENSE-2.0
  9#
 10# Unless required by applicable law or agreed to in writing, software
 11# distributed under the License is distributed on an "AS IS" BASIS,
 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 13# See the License for the specific language governing permissions and
 14# limitations under the License.
 15# ----------------------------------------------------------------------------
 16
 17from __future__ import annotations
 18
 19import torch
 20import numpy as np
 21import gymnasium as gym
 22
 23from embodichain.lab.gym.envs import BaseEnv, EnvCfg
 24from embodichain.lab.sim import SimulationManagerCfg
 25from embodichain.lab.visualization import VisualizationCfg
 26from embodichain.lab.sim.types import EnvAction, EnvObs
 27from embodichain.lab.sim.shapes import CubeCfg
 28from embodichain.lab.sim.objects import RigidObject, Robot
 29from embodichain.lab.sim.cfg import (
 30    RenderCfg,
 31    RobotCfg,
 32    RigidObjectCfg,
 33    RigidBodyAttributesCfg,
 34)
 35from embodichain.lab.gym.utils.registration import register_env
 36
 37
 38@register_env("RandomReach-v1", override=True)
 39class RandomReachEnv(BaseEnv):
 40
 41    robot_init_qpos = np.array(
 42        [1.57079, -1.57079, 1.57079, -1.57079, -1.57079, -3.14159]
 43    )
 44
 45    def __init__(
 46        self,
 47        num_envs=1,
 48        headless=False,
 49        device="cpu",
 50        renderer="hybrid",
 51        visualization: VisualizationCfg | None = None,
 52        **kwargs,
 53    ) -> None:
 54        env_cfg = EnvCfg(
 55            sim_cfg=SimulationManagerCfg(
 56                headless=headless,
 57                arena_space=2.0,
 58                sim_device=device,
 59                render_cfg=RenderCfg(renderer=renderer),
 60                visualization=visualization or VisualizationCfg(),
 61            ),
 62            num_envs=num_envs,
 63        )
 64
 65        super().__init__(
 66            cfg=env_cfg,
 67            **kwargs,
 68        )
 69
 70    def _setup_robot(self, **kwargs) -> Robot:
 71        from embodichain.data import get_data_path
 72
 73        file_path = get_data_path("UniversalRobots/UR10/UR10.urdf")
 74
 75        robot: Robot = self.sim.add_robot(
 76            cfg=RobotCfg(
 77                uid="ur10",
 78                fpath=file_path,
 79                init_pos=(0, 0, 1),
 80                init_qpos=self.robot_init_qpos,
 81            )
 82        )
 83
 84        qpos_limits = robot.body_data.qpos_limits[0].cpu().numpy()
 85        self.single_action_space = gym.spaces.Box(
 86            low=qpos_limits[:, 0], high=qpos_limits[:, 1], dtype=np.float32
 87        )
 88
 89        return robot
 90
 91    def _prepare_scene(self, **kwargs) -> None:
 92        size = 0.03
 93        # Create a kinematic cube object without collision.
 94        # Currently, we use this workaround for visualization purposes.
 95        self.cube: RigidObject = self.sim.add_rigid_object(
 96            cfg=RigidObjectCfg(
 97                uid="cube",
 98                shape=CubeCfg(size=[size, size, size]),
 99                attrs=RigidBodyAttributesCfg(enable_collision=False),
100                init_pos=(0.0, 0.0, 0.5),
101                body_type="kinematic",
102            ),
103        )
104
105    def _update_sim_state(self, **kwargs) -> None:
106        pose = torch.eye(4, device=self.device)
107        pose = pose.unsqueeze_(0).repeat(self.num_envs, 1, 1)
108        pose[:, :3, 3] += torch.rand(self.num_envs, 3, device=self.device) * 0.5 - 0.25
109        self.cube.set_local_pose(pose=pose)
110
111    def _step_action(self, action: EnvAction) -> EnvAction:
112        self.robot.set_qpos(qpos=action)
113        return action
114
115    def _extend_obs(self, obs: EnvObs, **kwargs) -> EnvObs:
116        # You can also use `cube = self.sim.get_rigid_object("cube")` to access obj.
117        # obs["cube_position"] = self.cube.get_local_pose()[:, :3]
118        return obs
119
120
121if __name__ == "__main__":
122    import argparse
123    import time
124
125    from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser
126    from embodichain.lab.visualization import visualization_cfg_from_args
127
128    parser = argparse.ArgumentParser(
129        description="Demo for running a random reach environment."
130    )
131    add_env_launcher_args_to_parser(parser)
132    args = parser.parse_args()
133
134    env = gym.make(
135        "RandomReach-v1",
136        num_envs=args.num_envs,
137        headless=args.headless,
138        device=args.device,
139        renderer=args.renderer,
140        visualization=visualization_cfg_from_args(args),
141    )
142
143    for episode in range(10):
144        print("Episode:", episode)
145        env.reset()
146        start_time = time.time()
147        total_steps = 0
148
149        for i in range(100):
150            action = env.action_space.sample()
151            action = torch.as_tensor(
152                action, dtype=torch.float32, device=env.get_wrapper_attr("device")
153            )
154
155            init_pose = env.unwrapped.robot_init_qpos
156            init_pose = (
157                torch.as_tensor(
158                    init_pose,
159                    dtype=torch.float32,
160                    device=env.get_wrapper_attr("device"),
161                )
162                .unsqueeze_(0)
163                .repeat(env.get_wrapper_attr("num_envs"), 1)
164            )
165            action = (
166                init_pose
167                + torch.rand_like(
168                    action, dtype=torch.float32, device=env.get_wrapper_attr("device")
169                )
170                * 0.2
171                - 0.1
172            )
173
174            obs, reward, done, truncated, info = env.step(action)
175            total_steps += env.get_wrapper_attr("num_envs")
176
177        end_time = time.time()
178        elapsed_time = end_time - start_time
179        if elapsed_time > 0:
180            fps = total_steps / elapsed_time
181            print(f"Total steps: {total_steps}")
182            print(f"Elapsed time: {elapsed_time:.2f} seconds")
183            print(f"FPS: {fps:.2f}")
184        else:
185            print("Elapsed time is too short to calculate FPS.")
186
187    env.close()

The Code Explained#

This tutorial demonstrates how to create a custom RL environment by inheriting from envs.BaseEnv. The environment implements a simple reach task where a robot arm tries to reach randomly positioned targets.

Environment Registration#

First, we register the environment with the Gymnasium registry using the utils.registration.register_env() decorator:

@register_env("RandomReach-v1", override=True)
class RandomReachEnv(BaseEnv):

The decorator parameters define:

  • Environment ID: "RandomReach-v1" - unique identifier for the environment

  • override: Whether to override existing environment with same ID

Environment Initialization#

The __init__ method configures the simulation environment and calls the parent constructor:

from embodichain.lab.visualization import VisualizationCfg
from embodichain.lab.sim.types import EnvAction, EnvObs
from embodichain.lab.sim.shapes import CubeCfg
from embodichain.lab.sim.objects import RigidObject, Robot
from embodichain.lab.sim.cfg import (
    RenderCfg,
    RobotCfg,
    RigidObjectCfg,
    RigidBodyAttributesCfg,
)
from embodichain.lab.gym.utils.registration import register_env


@register_env("RandomReach-v1", override=True)
class RandomReachEnv(BaseEnv):

    robot_init_qpos = np.array(
        [1.57079, -1.57079, 1.57079, -1.57079, -1.57079, -3.14159]
    )

    def __init__(
        self,

Key configuration options include:

  • num_envs: Number of parallel environments to run

  • headless: Whether to run without GUI (useful for training)

  • device: Computation device (“cpu” or “cuda”)

Robot Setup#

The _setup_robot method loads and configures the robot for the environment:

    def _setup_robot(self, **kwargs) -> Robot:
        from embodichain.data import get_data_path

        file_path = get_data_path("UniversalRobots/UR10/UR10.urdf")

        robot: Robot = self.sim.add_robot(
            cfg=RobotCfg(
                uid="ur10",
                fpath=file_path,
                init_pos=(0, 0, 1),
                init_qpos=self.robot_init_qpos,
            )
        )

        qpos_limits = robot.body_data.qpos_limits[0].cpu().numpy()
        self.single_action_space = gym.spaces.Box(
            low=qpos_limits[:, 0], high=qpos_limits[:, 1], dtype=np.float32
        )

        return robot

This method demonstrates:

  1. URDF Loading: Using data module to access robot URDF files

  2. Robot Configuration: Setting initial position and joint configuration

  3. Action Space Definition: Creating action space based on joint limits

The action space is automatically derived from the robot’s joint limits, ensuring actions stay within valid ranges.

Scene Preparation#

The _prepare_scene() method adds additional objects to the simulation environment:


        file_path = get_data_path("UniversalRobots/UR10/UR10.urdf")

        robot: Robot = self.sim.add_robot(
            cfg=RobotCfg(
                uid="ur10",
                fpath=file_path,
                init_pos=(0, 0, 1),
                init_qpos=self.robot_init_qpos,
            )
        )

        qpos_limits = robot.body_data.qpos_limits[0].cpu().numpy()

In this example, we add a kinematic cube that serves as a visual target. The cube is configured with:

  • No collision: enable_collision=False for visualization only

  • Kinematic body: Can be moved programmatically without physics

  • Custom size: Small 3cm cube for target visualization

  • initial position: Initially placed at a fixed location

State Updates#

The _update_sim_state method is called at each simulation step to update object states:

    def _update_sim_state(self, **kwargs) -> None:
        pose = torch.eye(4, device=self.device)
        pose = pose.unsqueeze_(0).repeat(self.num_envs, 1, 1)
        pose[:, :3, 3] += torch.rand(self.num_envs, 3, device=self.device) * 0.5 - 0.25
        self.cube.set_local_pose(pose=pose)

This method randomizes the cube’s position. The pose is updated for all parallel environments simultaneously.

Note that this method is called after perform action execution and simulation update but before observation collection. For more details, see envs.BaseEnv.step().

Action Execution#

The _step_action method applies actions to the robot:

    def _step_action(self, action: EnvAction) -> EnvAction:
        self.robot.set_qpos(qpos=action)
        return action

In this simple environment, actions directly set joint positions. More complex environments might:

  • Convert actions to joint torques or velocities

  • Apply action filtering or scaling

  • Implement inverse kinematics for end-effector control

Observation Extension#

The default observations include the following keys:

  • robot: Robot proprioception data (joint positions, velocities, efforts)

  • sensor (optional): Data from any sensors (e.g., cameras)

The _extend_obs method allows you to add custom observations:

    def _extend_obs(self, obs: EnvObs, **kwargs) -> EnvObs:
        # You can also use `cube = self.sim.get_rigid_object("cube")` to access obj.
        # obs["cube_position"] = self.cube.get_local_pose()[:, :3]
        return obs

While commented out in this example, you can add custom data like:

  • Object positions and orientations

  • Distance calculations

  • Custom sensor readings

  • Task-specific state information

The Code Execution#

To run the environment:

cd /path/to/embodichain
python scripts/tutorials/gym/random_reach.py

You can customize the execution with command-line options:

# Run multiple parallel environments
python scripts/tutorials/gym/random_reach.py --num_envs 4

# Run with GPU acceleration
python scripts/tutorials/gym/random_reach.py --device cuda

# Run in headless mode (no GUI)
python scripts/tutorials/gym/random_reach.py --headless

The script demonstrates:

  1. Environment Creation: Using gym.make() with custom parameters

  2. Episode Loop: Running multiple episodes with random actions

  3. Performance Monitoring: Calculating frames per second (FPS)

Key Features Demonstrated#

This tutorial showcases several important features of EmbodiChain environments:

  1. Gymnasium Integration: Full compatibility with the Gymnasium API

  2. Parallel Environments: Running multiple environments simultaneously for efficient training

  3. Robot Integration: Easy loading and control of robotic systems

  4. Custom Objects: Adding and manipulating scene objects

  5. Flexible Actions: Customizable action spaces and execution methods

  6. Extensible Observations: Adding task-specific observation data

Tip

Using an AI coding agent? Once you’re ready to create your own task environment, use the /add-task-env skill to scaffold the file with the correct structure, @register_env decorator, base class methods, and test stub. Use /add-test to write tests and /pre-commit-check to verify everything passes CI before committing.

Next Steps#

Contents