Data Generation#
This tutorial shows how to generate synthetic expert demonstration datasets using EmbodiChain’s built-in environment rollout and dataset manager. You will learn how to configure LeRobot recording in a gym config file (.json, .yaml, or .yml), how run_env.py builds an environment from configuration files, and how completed episodes are automatically saved to disk.
Overview#
EmbodiChain provides a built-in data generation workflow for imitation-learning and manipulation tasks:
Gym Configuration: Describes the scene, robot, sensors, randomization events, observations, dataset recorder, and rollout settings.
Action Configuration: Describes the task-specific expert action graph for tasks that use the action bank.
Environment Rollout: Builds the environment directly from configuration files and executes offline generation.
Expert Policy: Each task provides
create_demo_action_list()or another scripted policy entry to generate expert actions.Dataset Manager: Records observation-action pairs during
env.step().LeRobotRecorder: Converts completed episodes into LeRobot-compatible datasets, with optional video export.
What This Tutorial Records#
This page documents the full path from task configuration to saved dataset:
Prepare a task gym config (e.g.
gym_config.jsonorgym_config.yaml).Prepare an action config if the task uses the action bank (same supported extensions).
Launch the environment rollout with
run-env.Let the dataset manager automatically save completed episodes.
Example Task#
As a concrete example, this tutorial uses a real action-bank task shipped in the repository:
embodichain_tasks/configs/gym/pour_water/gym_config.jsondefines the simulation scene and dataset recording behavior (YAML equivalents such asembodichain_tasks/configs/gym/cobotmagic.yamlare also supported).embodichain_tasks/configs/gym/pour_water/action_config.jsondefines the action-bank graph used to solve the task.
The Code#
The tutorial corresponds to the run_env.py script in embodichain/lab/scripts.
Code for run_env.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
17import gymnasium
18import numpy as np
19import argparse
20import os
21import torch
22import tqdm
23
24from embodichain.lab.gym.utils.gym_utils import (
25 add_env_launcher_args_to_parser,
26 build_env_cfg_from_args,
27)
28from embodichain.lab.gym.utils.registration import (
29 discover_task_packages,
30 execute_init_hooks,
31)
32from embodichain.utils.logger import log_warning, log_info, log_error
33
34
35def generate_and_execute_action_list(env, idx, debug_mode, **kwargs):
36
37 action_list = env.get_wrapper_attr("create_demo_action_list")(
38 action_sentence=idx, **kwargs
39 )
40
41 if action_list is None or len(action_list) == 0:
42 log_warning("Action is invalid. Skip to next generation.")
43 return False
44
45 for action in tqdm.tqdm(
46 action_list, desc=f"Executing action list #{idx}", unit="step"
47 ):
48 # Step the environment with the current action
49 # The environment will automatically detect truncation based on action_length
50 obs, reward, terminated, truncated, info = env.step(action)
51
52 # TODO: We may assume in export demonstration rollout, there is no truncation from the env.
53 # but truncation is useful to improve the generation efficiency.
54
55 return True
56
57
58def generate_function(
59 env,
60 num_traj,
61 time_id: int = 0,
62 save_path: str = "",
63 save_video: bool = False,
64 debug_mode: bool = False,
65 **kwargs,
66):
67 """Generate and execute a sequence of actions in the environment.
68
69 This function resets the environment, generates and executes action trajectories,
70 collects data, and optionally saves videos of the episodes. It supports both online
71 and offline data generation modes.
72
73 Args:
74 env: The environment instance.
75 num_traj (int): Number of trajectories to generate per episode.
76 time_id (int, optional): Identifier for the current time step or episode.
77 save_path (str, optional): Path to save generated videos.
78 save_video (bool, optional): Whether to save episode videos.
79 debug_mode (bool, optional): Enable debug mode for visualization and logging.
80 **kwargs: Additional keyword arguments for data generation.
81
82 Returns:
83 bool: True if data generation is successful, False otherwise.
84 """
85
86 valid = True
87 _, _ = env.reset()
88 while True:
89 ret = []
90 for trajectory_idx in range(num_traj):
91 valid = generate_and_execute_action_list(
92 env, trajectory_idx, debug_mode, **kwargs
93 )
94
95 if not valid:
96 # Failed execution: reset without saving invalid data
97 _, _ = env.reset(options={"save_data": False})
98 break
99
100 if valid:
101 break
102 else:
103 log_warning("Reset valid flag to True.")
104 valid = True
105
106 return True
107
108
109def main(args, env, gym_config):
110 if getattr(args, "preview", False):
111 log_info(
112 "Preview mode enabled. Launching environment preview...", color="green"
113 )
114 preview(env)
115
116 log_info("Start offline data generation.", color="green")
117 # TODO: Support multiple trajectories per episode generation.
118 num_traj = 1
119 for i in range(gym_config.get("max_episodes", 1)):
120 generate_function(
121 env,
122 num_traj,
123 i,
124 save_path=getattr(args, "save_path", ""),
125 save_video=getattr(args, "save_video", False),
126 debug_mode=getattr(args, "debug_mode", False),
127 regenerate=getattr(args, "regenerate", False),
128 )
129
130 # Final reset.
131 _, _ = env.reset()
132
133
134def preview(env: gymnasium.Env) -> None:
135 """
136 Run the following code to create a demonstration and perform env steps.
137
138 ```
139 # Demo version of environment rollout
140 for i in range(10):
141 qpos = env.robot.get_qpos()
142
143 obs, reward, terminated, truncated, info = env.step(qpos)
144
145 # reset the environment
146 env.reset()
147 ```
148
149 Run the following code to preview the sensor observations.
150
151 ```
152 env.preview_sensor_data("camera")
153 ```
154 """
155 _, _ = env.reset()
156
157 end = False
158 while end is False:
159 print("Press `p` to enter embed mode to interact with the environment.")
160 print("Press `q` to quit the simulation.")
161 txt = input()
162 if txt == "p":
163 try:
164 from IPython import embed
165 except ImportError:
166 log_error(
167 "IPython is not installed. Preview mode requires IPython to be "
168 "available. Please install it with `pip install ipython` and try again."
169 )
170 continue
171
172 embed()
173 elif txt == "q":
174 end = True
175
176 exit(0)
177
178
179def cli():
180 """Command-line interface for environment runner.
181
182 Parses CLI arguments, builds the environment config, and launches
183 the data generation or preview workflow.
184 """
185 np.set_printoptions(5, suppress=True)
186 torch.set_printoptions(precision=5, sci_mode=False)
187
188 parser = argparse.ArgumentParser()
189
190 add_env_launcher_args_to_parser(parser)
191
192 args = parser.parse_args()
193
194 # Step 1: Discover all task packages via entry_points
195 discover_task_packages()
196
197 # Step 2: Execute init hooks (register managers, asset resolvers, etc.)
198 execute_init_hooks()
199
200 env_cfg, gym_config, action_config = build_env_cfg_from_args(args)
201
202 env = gymnasium.make(id=gym_config["id"], cfg=env_cfg, **action_config)
203
204 main(args, env, gym_config)
205
206
207if __name__ == "__main__":
208 cli()
The Code Explained#
The rollout script builds the environment from configuration, generates expert trajectories, executes them step by step, and relies on the dataset manager to auto-save valid episodes.
Step 1: Prepare the Task Configuration#
The first input to the pipeline is the task gym config file. In the example below, the same file contains rollout settings, scene randomization, observations, dataset recording, and robot or sensor definitions.
The rollout settings include the episode count:
"id": "PourWater-v3",
"max_episodes": 5,
"max_episode_steps": 300,
The dataset-related part looks like this:
"quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635]
}
},
{
"sensor_type": "Camera",
"uid": "cam_left_wrist",
"width": 640,
"height": 480,
"intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812],
"extrinsics": {
"parent": "left_link6",
"pos": [-0.08, 0.0, 0.04],
"quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635]
}
}
],
"light": {
"direct": [
{
"uid": "light_1",
"light_type": "point",
Important parameters are:
max_episodes: Number of rollout episodes generated by
run_env.py.max_episode_steps: Maximum number of environment steps per episode.
dataset.lerobot.params.robot_meta: Robot metadata such as robot type and control frequency.
dataset.lerobot.params.instruction: Task language instruction stored together with the dataset.
dataset.lerobot.params.extra: Additional metadata such as scene type and task description.
dataset.lerobot.params.use_videos: Whether camera observations should be stored as videos.
env.control_parts: Controlled robot parts in the environment.
In the current implementation, LeRobotRecorder stores robot state and action features following LeRobot official format: observation.state for joint positions, action for applied actions, and observation.images.{sensor_name} for camera images.
Step 2: Prepare the Action Configuration#
For tasks that use the action bank, the second input is action_config.json. This file defines the expert action graph consumed by create_demo_action_list(). In the example below, the file is organized around scope, node, edge, and sync.
Action bank structure in the example task Pour_Water
Scope Configuration
"scope": {
"right_arm": {
"type": "DiGraph",
"dim": [
6
],
"init": {
"method": "current_qpos",
"init_node_name": "right_arm_init_qpos"
},
"dtype": "float32"
},
"left_arm": {
"type": "DiGraph",
"dim": [
6
],
"init": {
"method": "current_qpos",
"init_node_name": "left_arm_init_qpos"
},
"dtype": "float32"
},
"left_eef": {
"type": "DiGraph",
"dim": [
1
],
"init": {
"method": "given_qpos",
"kwargs": {
"given_qpos": [
1
]
},
"init_node_name": ""
},
"dtype": "float32"
},
"right_eef": {
"type": "DiGraph",
"dim": [
1
],
"init": {
"method": "given_qpos",
"kwargs": {
"given_qpos": [
1
]
},
"init_node_name": ""
},
"dtype": "float32"
}
},
Node Configuration
"bottle_grasp": {
"name": "generate_affordances_from_src",
"kwargs": {
"affordance_infos": [
{
"src_key": "bottle_pose",
"dst_key": "bottle_grasp_pose",
"valid_funcs_name_kwargs_proc": [
{
"name": "no_validation",
"kwargs": {},
"pass_processes": [
{
"name": "get_rotation_replaced_pose",
"kwargs": {
"rotation_value": "env.affordance_datas['right_arm_aim_qpos'][0]",
"rot_axis": "z",
"mode": "intrinsic"
}
},
{
"name": "get_frame_changed_pose",
"kwargs": {
"frame_change_matrix": "env.affordance_datas['bottle_pose']",
"mode": "intrinsic",
"inverse": true
}
},
{
"name": "get_frame_changed_pose",
"kwargs": {
"frame_change_matrix": "env.affordance_datas['bottle_grasp_pose']",
"mode": "intrinsic"
}
}
]
}
]
},
{
"src_key": "bottle_grasp_pose",
"dst_key": "bottle_pre1_pose",
"valid_funcs_name_kwargs_proc": [
{
"name": "no_validation",
"kwargs": {},
"pass_processes": [
{
"name": "get_offset_pose",
"kwargs": {
"offset_value": -0.05,
"direction": "z",
"mode": "intrinsic"
}
}
]
}
]
},
{
"src_key": "bottle_pre1_pose",
"dst_key": "bottle_pre2_pose",
"valid_funcs_name_kwargs_proc": [
{
"name": "no_validation",
"kwargs": {},
"pass_processes": [
{
"name": "get_offset_pose",
"kwargs": {
"offset_value": -0.05,
"direction": "z",
"mode": "intrinsic"
}
}
]
}
]
}
]
}
}
Edge Configuration
"up_to_move": {
"src": "bottle_up_qpos",
"sink": "pour_water_start_qpos",
"duration": 20,
"name": "plan_trajectory",
"kwargs": {
"agent_uid": "right_arm",
"keypose_names": [
"bottle_up_qpos",
"pour_water_start_qpos"
]
}
}
},
{
"move_to_rotation": {
"src": "pour_water_start_qpos",
"sink": "bottle_rotation_qpos",
"duration": 24,
"name": "plan_trajectory",
"kwargs": {
"agent_uid": "right_arm",
"keypose_names": [
"pour_water_start_qpos",
"bottle_rotation_qpos"
]
}
}
Synchronization
"sync": {
"rclose0": {
"depend_tasks": [
"pre1_to_grasp"
]
},
"grasp_to_up": {
"depend_tasks": [
"rclose0"
]
},
"ropen0": {
"depend_tasks": [
"pre_place_back_to_place"
]
},
"place_back_to_init": {
"depend_tasks": [
"ropen0"
]
},
"left_arm_go_back": {
"depend_tasks": [
"ropen0"
]
}
},
This structure defines the expert rollout as follows:
Scope: Defines controllable sub-graphs such as
right_arm,left_arm,right_eef, andleft_eef.Node: Defines key poses, targets computed from object affordances, and IK-generated joint targets.
Edge: Defines executable transitions between nodes, including duration and execution function.
Sync: Defines execution order rules between independently configured sub-actions.
Note: Action bank is not the only way to generate demonstrations. Depending on the task design, trajectories can also be produced by other scripted generation methods.
Step 3: Launch the Environment Rollout#
The rollout script parses command-line arguments, loads the gym and action config files, converts them into environment configuration objects, creates the environment instance, and then runs offline rollout for max_episodes episodes:
def cli():
"""Command-line interface for environment runner.
Parses CLI arguments, builds the environment config, and launches
the data generation or preview workflow.
"""
np.set_printoptions(5, suppress=True)
torch.set_printoptions(precision=5, sci_mode=False)
parser = argparse.ArgumentParser()
add_env_launcher_args_to_parser(parser)
args = parser.parse_args()
# Step 1: Discover all task packages via entry_points
discover_task_packages()
# Step 2: Execute init hooks (register managers, asset resolvers, etc.)
execute_init_hooks()
env_cfg, gym_config, action_config = build_env_cfg_from_args(args)
env = gymnasium.make(id=gym_config["id"], cfg=env_cfg, **action_config)
main(args, env, gym_config)
Each rollout internally calls create_demo_action_list(), validates the returned sequence, executes actions with env.step(action), and discards invalid rollouts by resetting with save_data=False.
The recommended CLI entrypoint is:
python -m embodichain run-env \
--gym_config embodichain_tasks/configs/gym/pour_water/gym_config.json \
--action_config embodichain_tasks/configs/gym/pour_water/action_config.json \
--headless
For interactive inspection, you can use preview mode: replace --headless with --preview.
When --preview is enabled, the script opens the environment in an interactive debugging mode. This mode is for inspection and does not save datasets.
Useful CLI arguments:
–gym_config: Path to the task config file (
.json,.yaml, or.yml).–action_config: Path to the action-bank config file (
.json,.yaml, or.yml).–num_envs: Number of environments to run in parallel.
–device: Simulation device, such as
cpuorcuda.–headless: Run without GUI for faster generation.
–enable_rt: Enable ray tracing for higher-quality visual observations.
–preview: Launch the environment in interactive preview mode.
–filter_dataset_saving: Disable dataset saving for debugging.
For the complete CLI argument list, see CLI Reference.
Outputs#
After successful execution, completed episodes are saved under the configured dataset root. A LeRobot dataset typically contains:
If no explicit save path is provided and EMBODICHAIN_DATASET_ROOT is not set, LeRobotRecorder uses ~/.cache/embodichain_datasets as the default dataset root.
data/: Recorded action and state data.
videos/: Camera observations saved as videos when
use_videos=True.meta/: Dataset metadata such as task information and robot description.
Dataset folders are automatically numbered, which makes it easy to run repeated generations without overwriting previous results.
In a practical workflow, the output of this stage is the synthesized dataset itself. Later training scripts typically consume these saved LeRobot episodes instead of regenerating trajectories each time.
Best Practices#
Keep the config pair together: Version gym and action configs together for action-bank tasks (either JSON or YAML).
Use valid scripted policies: Make sure
create_demo_action_list()returns executable trajectories for the current scene.Use ``–headless`` for throughput: Disable the GUI when generating large datasets.
Use ``–preview`` and ``–filter_dataset_saving`` for debugging: Inspect task logic without writing datasets.
Discard invalid rollouts: Keep the default validation logic so failed trajectories are not saved.