castanetnicolas/robosuite_manipulation

Robosuite manipulation (LeRobot EnvHub)

A self-contained LeRobot EnvHub environment: a robot arm performing manipulation tasks, built on robosuite.

Robots

The arm is selectable via robot=, the same way the problem is selectable via task=:

robot	robosuite arm	gripper	agent_pos
ur5e	UR5e	RethinkGripper	9-D
panda	Panda (Franka Emika)	PandaGripper	9-D

ur5e is the default. It is pinned to the RethinkGripper rather than its robosuite default (the Robotiq85), whose curling fingers cannot grasp the flat nut handle. The Panda keeps its own gripper. Both pair with a 2-DoF gripper qpos, so agent_pos is 9-D either way and a dataset recorded on one is schema-compatible with the other.

Note the bundled scripted experts are tuned against the RethinkGripper's geometry on the UR5e and refuse to run on another robot rather than silently emitting unusable demonstrations.

Tasks

It exposes two robosuite tasks (single-object mode), selectable via task=:

• pick-and-place (default, pick_place_can) — grasp one object and drop it in the correct bin;
• nut assembly (nut_assembly_round / nut_assembly_square) — grasp a nut and fit it onto its peg.

Quick start

Once published to the Hub, load it in one line:

from lerobot.envs import make_env

# trust_remote_code=True is required to execute the hub env.py
envs = make_env("castanetnicolas/robosuite_manipulation", n_envs=4, trust_remote_code=True)

suite_name = next(iter(envs))
env = envs[suite_name][0]
obs, info = env.reset(seed=0)
action = env.action_space.sample()
obs, reward, terminated, truncated, info = env.step(action)
env.close()

Requirements

pip install -r requirements.txt

robosuite depends on MuJoCo, which needs a working (headless-capable) OpenGL stack for offscreen rendering. On a server, set e.g. MUJOCO_GL=egl.

API

The package exposes the single EnvHub entry point:

def make_env(n_envs: int = 1, use_async_envs: bool = False, cfg=None) -> gym.vector.VectorEnv

• n_envs — number of parallel rollouts.
• use_async_envs — AsyncVectorEnv (one MuJoCo process per env) vs SyncVectorEnv.
• cfg — optional config object; honored attributes: task, robot, obs_type, observation_width, observation_height, episode_length, fps (→ control_freq), and a gym_kwargs dict that overrides the rest.

Tasks

Two robosuite environments are exposed, both in single-object mode (single_object_mode=2). Pick-and-place (base PickPlace, object selected by object_type) — drop the object in the correct bin:

task id	object	object_type
pick_place_can	can	can
pick_place_milk	milk	milk
pick_place_bread	bread	bread
pick_place_cereal	cereal	cereal

Nut assembly (base NutAssembly, nut selected by nut_type) — fit the nut onto its peg:

task id	nut	nut_type
nut_assembly_round	round	round
nut_assembly_square	square	square

The action space (7-D OSC), state (9-D) and cameras are identical across all tasks and both robots. Note: the bundled scripted expert (generate_pick_and_place_trajs.py) only solves the pick_place_* tasks; nut assembly needs a peg-insertion controller.

Observation & action spaces

Observations are returned in LeRobot format (obs_type="pixels_agent_pos", the default):

• pixels — a dict keyed by camera name, each an (H, W, 3) uint8 RGB image (one camera by default, agentview). LeRobot maps these to observation.images.<cam>, the layout image/state policies (e.g. pi0) expect. For multiple cameras pass a comma-separated camera_name, e.g. "agentview,robot0_eye_in_hand".
• agent_pos — end-effector proprioception [eef_pos(3), eef_quat(4), gripper_qpos(2)] (9-D, RethinkGripper), mapped to observation.state.

Other modes: pixels (images only) and state (flat proprio + object vector via robosuite's GymWrapper; note this mode is not compatible with lerobot-eval, which expects a dict observation).

The action is the robosuite OSC end-effector command: a 6-D pose delta plus a 1-D gripper command, bounded by the controller's action_spec.

Reward & termination

• Reward — dense shaped reward from robosuite (reward_shaping=True).
• terminated — the object is placed in the correct bin (_check_success()).
• truncated — the episode reaches episode_length steps (default 500).

Local testing

python env.py            # runs a short random-action smoke test

Or, following the EnvHub docs:

from lerobot.envs.utils import _load_module_from_path, _call_make_env, _normalize_hub_result

module = _load_module_from_path("./env.py")
result = _call_make_env(module, n_envs=2, use_async_envs=False, cfg=None)
normalized = _normalize_hub_result(result)
env = normalized[next(iter(normalized))][0]
obs, info = env.reset()
env.close()

Local development

These helper scripts live alongside env.py but import lerobot, so they are not part of the portable EnvHub package. When uploading to the Hub, publish only env.py, requirements.txt, and this README.md.

file	what it does
test_env_local.py	The EnvHub "test locally" flow (load env.py → make_env → reset).
rl_state_demo.py	Runs LeRobot's SAC actor (gaussian_actor) on the state obs (integration smoke test, untrained).
inspect_pi0_action_space.py	Dumps a pretrained checkpoint's action/state normalization stats (e.g. lerobot/pi0) to compare its action scale against this env's.

For evaluating a policy against this env from the repo checkout, a first-class EnvConfig is registered as --env.type=robosuite (see lerobot/envs/configs.py), so you can run e.g.:

lerobot-eval --policy.path=<checkpoint> --policy.device=cuda \
  --env.type=robosuite --eval.n_episodes=20 --eval.batch_size=4