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:
configs/gym/pour_water/gym_config.jsondefines the simulation scene and dataset recording behavior (YAML equivalents such asconfigs/gym/cobotmagic.yamlare also supported).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.utils.logger import log_warning, log_info, log_error
29
30
31def generate_and_execute_action_list(env, idx, debug_mode, **kwargs):
32
33 action_list = env.get_wrapper_attr("create_demo_action_list")(
34 action_sentence=idx, **kwargs
35 )
36
37 if action_list is None or len(action_list) == 0:
38 log_warning("Action is invalid. Skip to next generation.")
39 return False
40
41 for action in tqdm.tqdm(
42 action_list, desc=f"Executing action list #{idx}", unit="step"
43 ):
44 # Step the environment with the current action
45 # The environment will automatically detect truncation based on action_length
46 obs, reward, terminated, truncated, info = env.step(action)
47
48 # TODO: We may assume in export demonstration rollout, there is no truncation from the env.
49 # but truncation is useful to improve the generation efficiency.
50
51 return True
52
53
54def generate_function(
55 env,
56 num_traj,
57 time_id: int = 0,
58 save_path: str = "",
59 save_video: bool = False,
60 debug_mode: bool = False,
61 **kwargs,
62):
63 """Generate and execute a sequence of actions in the environment.
64
65 This function resets the environment, generates and executes action trajectories,
66 collects data, and optionally saves videos of the episodes. It supports both online
67 and offline data generation modes.
68
69 Args:
70 env: The environment instance.
71 num_traj (int): Number of trajectories to generate per episode.
72 time_id (int, optional): Identifier for the current time step or episode.
73 save_path (str, optional): Path to save generated videos.
74 save_video (bool, optional): Whether to save episode videos.
75 debug_mode (bool, optional): Enable debug mode for visualization and logging.
76 **kwargs: Additional keyword arguments for data generation.
77
78 Returns:
79 bool: True if data generation is successful, False otherwise.
80 """
81
82 valid = True
83 _, _ = env.reset()
84 while True:
85 ret = []
86 for trajectory_idx in range(num_traj):
87 valid = generate_and_execute_action_list(
88 env, trajectory_idx, debug_mode, **kwargs
89 )
90
91 if not valid:
92 # Failed execution: reset without saving invalid data
93 _, _ = env.reset(options={"save_data": False})
94 break
95
96 if valid:
97 break
98 else:
99 log_warning("Reset valid flag to True.")
100 valid = True
101
102 return True
103
104
105def main(args, env, gym_config):
106 if getattr(args, "preview", False):
107 log_info(
108 "Preview mode enabled. Launching environment preview...", color="green"
109 )
110 preview(env)
111
112 log_info("Start offline data generation.", color="green")
113 # TODO: Support multiple trajectories per episode generation.
114 num_traj = 1
115 for i in range(gym_config.get("max_episodes", 1)):
116 generate_function(
117 env,
118 num_traj,
119 i,
120 save_path=getattr(args, "save_path", ""),
121 save_video=getattr(args, "save_video", False),
122 debug_mode=getattr(args, "debug_mode", False),
123 regenerate=getattr(args, "regenerate", False),
124 )
125
126 # Final reset.
127 _, _ = env.reset()
128
129
130def preview(env: gymnasium.Env) -> None:
131 """
132 Run the following code to create a demonstration and perform env steps.
133
134 ```
135 # Demo version of environment rollout
136 for i in range(10):
137 qpos = env.robot.get_qpos()
138
139 obs, reward, terminated, truncated, info = env.step(qpos)
140
141 # reset the environment
142 env.reset()
143 ```
144
145 Run the following code to preview the sensor observations.
146
147 ```
148 env.preview_sensor_data("camera")
149 ```
150 """
151 _, _ = env.reset()
152
153 end = False
154 while end is False:
155 print("Press `p` to enter embed mode to interact with the environment.")
156 print("Press `q` to quit the simulation.")
157 txt = input()
158 if txt == "p":
159 try:
160 from IPython import embed
161 except ImportError:
162 log_error(
163 "IPython is not installed. Preview mode requires IPython to be "
164 "available. Please install it with `pip install ipython` and try again."
165 )
166 continue
167
168 embed()
169 elif txt == "q":
170 end = True
171
172 exit(0)
173
174
175def cli():
176 """Command-line interface for environment runner.
177
178 Parses CLI arguments, builds the environment config, and launches
179 the data generation or preview workflow.
180 """
181 np.set_printoptions(5, suppress=True)
182 torch.set_printoptions(precision=5, sci_mode=False)
183
184 parser = argparse.ArgumentParser()
185
186 add_env_launcher_args_to_parser(parser)
187
188 args = parser.parse_args()
189
190 env_cfg, gym_config, action_config = build_env_cfg_from_args(args)
191
192 env = gymnasium.make(id=gym_config["id"], cfg=env_cfg, **action_config)
193
194 main(args, env, gym_config)
195
196
197if __name__ == "__main__":
198 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()
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 configs/gym/pour_water/gym_config.json \
--action_config 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.