Rate this Page

★ ★ ★ ★ ★

tutorials/mujoco_cube_bowl_macros

Run in Google Colab

Note

Go to the end to download the full example code.

MuJoCo scripted manipulation with human-readable robot actions#

The first step in controlling a complex robot is rarely an end-to-end RL policy. Closed-loop controllers, scripted motion primitives, and imitation data are often used first to make the task reliable enough for later reward-driven fine tuning. This tutorial uses a MuJoCo Menagerie UR5e arm with a Robotiq 2F-85 gripper and shows how to ask the robot to do human-readable things such as “reach this pose”, “close the gripper”, or “go home” while TorchRL handles the low-level action plumbing.

We will control a cube-to-bowl task without learning. A policy emits a single RobotMacroAction object under the normal "action" key, and a TorchRL transform expands and executes that high-level command through MuJoCo. The same pattern can later provide demonstrations, curricula, or a safe initialization for residual RL.

What you will learn#

how to instantiate a Menagerie-backed custom MuJoCo environment;
how to use RobotMacroAction as a readable action object;
how URScriptPrimitiveTransform lets env.step(td) consume those actions directly;
how to hold the end effector still for gripper-only commands and keep its orientation constrained along a Cartesian motion;
how to write a scripted contact-rich cube-to-bowl policy as an explicit list of poses and gripper commands;
how to request rendered pixels from the environment and record them with VideoRecorder;
how to reset and reuse one environment across scripted scenes.

from __future__ import annotations

import importlib.util
import os

from collections.abc import Iterator

import torch
from tensordict import TensorDict, TensorDictBase
from torchrl.envs import CubeBowlEnv, RobotMacroAction, RobotMacroActionMode, step_mdp
from torchrl.record import VideoRecorder

_has_mujoco = importlib.util.find_spec("mujoco") is not None

if not _has_mujoco:
    raise ImportError("This tutorial requires the `mujoco` Python package.")

Environment#

The scene uses the MuJoCo Menagerie UR5e and Robotiq 2F-85 models. Set TORCHRL_MUJOCO_MENAGERIE_PATH to a checkout containing the universal_robots_ur5e and robotiq_2f85 assets:

git clone --depth=1 --filter=blob:none --sparse \
    https://github.com/google-deepmind/mujoco_menagerie.git /tmp/menagerie
git -C /tmp/menagerie sparse-checkout set \
    universal_robots_ur5e robotiq_2f85
export TORCHRL_MUJOCO_MENAGERIE_PATH=/tmp/menagerie

CubeBowlEnv exposes the state needed for scripted manipulation: robot joints, gripper joints, the pinch site, the cube pose, the bowl target and a success flag. The task reward is intentionally sparse. It is 1 only when the cube center is within the environment’s placement tolerance of the bowl target coordinate, and 0 otherwise.

MENAGERIE_PATH = os.environ.get(CubeBowlEnv.MENAGERIE_ENV_VAR)
if MENAGERIE_PATH is None:
    raise RuntimeError(
        f"Set {CubeBowlEnv.MENAGERIE_ENV_VAR} to a MuJoCo Menagerie checkout "
        "before running this tutorial."
    )

CI fast path#

Rendering dominates the runtime of this tutorial, and the doc build does not need four randomized rollouts to illustrate the API. When TORCHRL_TUTORIALS_FAST=1 (set in the docs CI workflow), we render at a smaller resolution and run a single randomized rollout. Local builds keep the full-quality settings.

TUTORIAL_FAST = os.environ.get("TORCHRL_TUTORIALS_FAST", "0") == "1"

MAX_EPISODE_STEPS = 12000
RENDER_WIDTH = 320 if TUTORIAL_FAST else 480
RENDER_HEIGHT = 240 if TUTORIAL_FAST else 360
VIDEO_FRAME_SKIP = 16
VIDEO_INTERVAL_MS = 55
IK_KWARGS = {
    "iterations": 220,
    "orientation_weight": 1.0,
    "step_size": 0.7,
    "damping": 1e-4,
}

RobotMacroAction as the action#

A readable manipulation policy should not have to emit magic flat tensors or loose collections of root TensorDict keys. It can emit one structured object under "action". RobotMacroAction is that object:

RobotMacroAction.reach_pose(position=..., quaternion=...) moves the arm;
RobotMacroAction.close_gripper() closes only the gripper and keeps the arm where it is;
RobotMacroAction.reach_pose(..., gripper="closed") moves the arm and closes the gripper in one macro action;
RobotMacroAction.wait() holds the complete low-level target, while open_gripper and close_gripper preserve the six arm-joint targets;
RobotMacroAction.RESET asks the environment to return to its configured home joint posture.

The symbolic gripper commands "open", "closed" and "keep" are resolved by the transform. User code never has to know the Robotiq actuator’s numeric command range. When we want a task-specific grasp, we can still make the target explicit in object units: close about one millimeter past the cube width, ask the environment to convert that width to a gripper command, and pass that command to the RobotMacroAction.

Create the MuJoCo scene#

We start by constructing the cube-and-bowl task once. from_pixels=True asks the environment to include rendered frames in the TensorDict, which a standard VideoRecorder can capture. We then append one primitive-control transform with execute=True. After that, env.step accepts a RobotMacroAction directly under td["action"] and executes the expanded low-level sequence internally.

# ``render_every`` matches the recorder's ``skip``: the recorder keeps one
# frame in ``VIDEO_FRAME_SKIP``, so rendering the frames it drops would only
# burn time (rendering dominates the cost of a step in this scene).
env = CubeBowlEnv(
    seed=0,
    max_episode_steps=MAX_EPISODE_STEPS,
    menagerie_path=MENAGERIE_PATH,
    from_pixels=True,
    pixels_only=False,
    render_width=RENDER_WIDTH,
    render_height=RENDER_HEIGHT,
    render_every=VIDEO_FRAME_SKIP,
)
assert env.action_spec.shape[-1] == 7
assert env.robot_home_qpos is not None

primitive_control = env.make_urscript_transform(
    execute=True,
    macro_steps=250,
    settle_steps=1600,
    ik_kwargs=IK_KWARGS,
    stack_rewards=True,
    stack_observations=False,
)
recorder = VideoRecorder(
    logger=None,
    tag="scripted_macro",
    skip=VIDEO_FRAME_SKIP,
    make_grid=False,
)
env = env.append_transform(recorder)
env = env.append_transform(primitive_control)

Indexing MuJoCo batches#

Native MuJoCo environment batches use snapshot indexing rather than the live worker views used by ParallelEnv. If a custom MuJoCo environment is created with multiple physics states, env[0] returns a detached one-element snapshot: stepping or resetting the snapshot does not mutate the parent batch until the snapshot is assigned back explicitly. Slices, integer NumPy arrays and integer torch tensors select larger snapshots, while boolean masks are intentionally not supported.

batched_env = CubeBowlEnv(..., num_envs=4)
env0 = batched_env[0]
env0.rand_step()
batched_env[0] = env0

Preserving a singleton batch for integer indexing keeps downstream code that expects batched specs and TensorDicts working unchanged.

cube_width = 2 * env.OBJECT_HALF_SIZE
grasp_width = cube_width - 0.001
gripper_close_command = env.gripper_ctrl_for_width(grasp_width)

workspace_layouts = torch.tensor(
    [
        # Cube in front of the robot, bowl behind it and slightly to the right.
        [[0.44, -0.21], [0.48, 0.10]],
        # Cube on the robot's right-front side, bowl on the left-back side.
        [[0.51, -0.18], [0.39, 0.04]],
        # Cube in front of the robot, bowl on the left-back side.
        [[0.45, -0.20], [0.38, 0.02]],
        # Cube on the robot's right-front side, bowl near the center-back.
        [[0.52, -0.18], [0.45, 0.02]],
    ],
)


def random_scene_reset(rollout_index: int) -> TensorDict:
    layout = workspace_layouts[rollout_index % len(workspace_layouts)]
    cube_xy = layout[0].view(1, 2)
    bowl_xy = layout[1].view(1, 2)
    cube_pos = torch.cat(
        [cube_xy, torch.full_like(cube_xy[..., :1], env.cube_position[2])],
        dim=-1,
    )
    bowl_target_z = env.bowl_position[2] + env.MENAGERIE_BOWL_TARGET_OFFSET[2]
    bowl_pos = torch.cat(
        [bowl_xy, torch.full_like(bowl_xy[..., :1], bowl_target_z)], -1
    )
    return TensorDict(
        {
            "cube_pos": cube_pos.to(dtype=env.dtype, device=env.device),
            "bowl_pos": bowl_pos.to(dtype=env.dtype, device=env.device),
        },
        batch_size=[1],
        device=env.device,
    )


td = env.reset()
assert td["robot_qpos"].shape[-1] == 6
assert env.low_level_action(td["robot_qpos"]).shape[-1] == 7
assert td["pixels"].shape[-3:] == torch.Size([RENDER_HEIGHT, RENDER_WIDTH, 3])

Policy shape#

A scripted policy can be just a generator of explicit robot commands. This is the minimal shape we want users to see: build a list of poses and gripper commands, place the next RobotMacroAction under the normal action key, and call env.step.

def gen_small_actions(td: TensorDictBase) -> Iterator[RobotMacroAction]:
    joints = td["robot_qpos"].clone()
    joints[..., 0] = joints[..., 0] + 0.35
    yield RobotMacroAction.reach_joints(joints=joints, gripper="open", steps=18)
    yield RobotMacroAction.close_gripper(steps=18)


small_actions = gen_small_actions(td)
small_action = next(small_actions)
td["action"] = small_action
assert td["action"].mode.shape == torch.Size([1, 1])

Keeping the end effector constrained#

There are two different ways to keep the end effector in place. When only the gripper should move, open_gripper() and close_gripper() copy the current arm target and modify only the gripper command. A wait() action holds both the arm and the gripper target.

During a Cartesian move, keeping the tool orientation fixed requires an orientation constraint. A full target quaternion locks all three rotational degrees of freedom. Often two are enough: for a level gripper, rotations about the world x and y axes must remain constrained, while spin about world z can stay free. orientation_mask=(1.0, 1.0, 0.0) expresses exactly that constraint and avoids over-constraining the inverse-kinematics solve.

By default, reach_pose solves inverse kinematics at the endpoint and then interpolates the arm joints. Both endpoint poses can be valid even though the tool tilts or leaves the straight Cartesian line between them. Setting path="cartesian" asks the solver to warm-start a new IK solve at every waypoint, so the position and orientation constraints hold throughout the motion. The carry actions below use both options:

RobotMacroAction.reach_pose(
    position=target_position,
    quaternion=level_quaternion,
    orientation_mask=(1.0, 1.0, 0.0),
    path="cartesian",
    gripper="closed",
)

Cartesian waypoint solving is more expensive than joint interpolation. It is most useful when intermediate geometry matters, such as transporting a held object without tilting it or drifting sideways.

A scripted cube-to-bowl macro#

We can now write the cube-to-bowl controller as a readable sequence of manual robot commands. The generator below is deliberately explicit: it names the poses we want the pinch site to reach and interleaves them with gripper commands.

home_qpos = torch.as_tensor(
    env.robot_home_qpos,
    dtype=td["robot_qpos"].dtype,
    device=td["robot_qpos"].device,
).view(1, 6)
policy_state: dict[str, TensorDictBase | torch.Tensor] = {
    "td": td,
    "gripper_quat": td["pinch_quat"].clone(),
}


def pose_at(
    xyz: torch.Tensor, quat: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor]:
    if quat is None:
        quat = torch.zeros(xyz.shape[:-1] + (4,), dtype=xyz.dtype, device=xyz.device)
        quat[..., 0] = 1.0
    else:
        quat = quat.to(dtype=xyz.dtype, device=xyz.device).expand(xyz.shape[:-1] + (4,))
    return xyz, quat


def carry_action(alpha: float, bowl: torch.Tensor) -> RobotMacroAction:
    current_td = policy_state["td"]
    gripper_quat = policy_state["gripper_quat"]
    cube = current_td["cube_pos"].clone()
    carry_start_cube = policy_state["carry_start_cube"]
    pinch_to_cube = current_td["pinch_pos"].clone() - cube
    desired_xy = carry_start_cube[..., :2] + alpha * (
        bowl[..., :2] - carry_start_cube[..., :2]
    )
    desired_cube = torch.cat(
        [
            desired_xy,
            torch.full_like(bowl[..., 2:3], 0.24),
        ],
        dim=-1,
    )
    position, quaternion = pose_at(desired_cube + pinch_to_cube, gripper_quat)
    return RobotMacroAction.reach_pose(
        position=position,
        quaternion=quaternion,
        orientation_mask=(1.0, 1.0, 0.0),
        path="cartesian",
        gripper="closed",
        gripper_command=gripper_close_command,
        steps=80,
        settle_steps=20,
    )


def gen_actions(td: TensorDictBase) -> Iterator:
    cube = td["cube_pos"].clone()
    bowl = td["bowl_pos"].clone()
    quat = td["pinch_quat"].clone()
    grasp_offset = cube.new_tensor([[0.0, 0.0, -0.016]])

    # Action 0: Keep the arm at the reset joint target and let the scene settle.
    yield RobotMacroAction.wait(gripper="open", steps=1, settle_steps=19)

    # Action 1: Fully open the gripper before approaching the cube.
    yield RobotMacroAction.open_gripper(steps=100, settle_steps=20)

    current_td = policy_state["td"]
    cube = current_td["cube_pos"].clone()
    bowl = current_td["bowl_pos"].clone()
    quat = current_td["pinch_quat"].clone()
    position, quaternion = pose_at(cube + cube.new_tensor([[0.0, 0.0, 0.18]]), quat)
    # Action 2: Move the open gripper above the cube.
    yield RobotMacroAction.reach_pose(
        position=position,
        quaternion=quaternion,
        gripper="open",
        steps=180,
        settle_steps=60,
    )

    position, quaternion = pose_at(cube + grasp_offset, quat)
    # Action 3: Lower the open gripper to the grasp pose around the cube.
    yield RobotMacroAction.reach_pose(
        position=position,
        quaternion=quaternion,
        gripper="open",
        steps=180,
        settle_steps=60,
    )

    # Action 4: Close the gripper to a grasp about one millimeter tighter than
    # the cube width.
    yield RobotMacroAction.close_gripper(
        command=gripper_close_command, steps=160, settle_steps=80
    )

    current_td = policy_state["td"]
    cube = current_td["cube_pos"].clone()
    pinch_to_cube = current_td["pinch_pos"].clone() - cube
    position, quaternion = pose_at(
        cube + pinch_to_cube + cube.new_tensor([[0.0, 0.0, 0.20]]),
        quat,
    )
    # Action 5: Lift the grasped cube above the table.
    yield RobotMacroAction.reach_pose(
        position=position,
        quaternion=quaternion,
        gripper="closed",
        gripper_command=gripper_close_command,
        steps=120,
        settle_steps=60,
    )

    # The cube is not welded to the gripper. It is held by contact and
    # friction between the pads. If we ask for one large move directly to the
    # bowl, the fingers can drag faster than the cube can safely follow: the
    # cube may slide inside the grasp, get pushed out, or arrive shifted enough
    # that the final drop misses the bowl. We therefore carry it through a few
    # intermediate macro targets in the table plane. Within each macro,
    # ``path="cartesian"`` keeps the pinch site on the straight Cartesian line
    # and ``orientation_mask`` keeps the gripper level. After each target, the
    # next command uses the latest observed cube and finger positions.
    policy_state["carry_start_cube"] = policy_state["td"]["cube_pos"].clone()

    # Action 6: Carry the cube one quarter of the way toward the bowl.
    yield carry_action(0.25, bowl)

    # Action 7: Carry the cube halfway toward the bowl.
    yield carry_action(0.50, bowl)

    # Action 8: Carry the cube three quarters of the way toward the bowl.
    yield carry_action(0.75, bowl)

    # Action 9: Carry the cube above the bowl center.
    yield carry_action(1.00, bowl)

    drop_cube = torch.cat(
        [
            bowl[..., :1],
            bowl[..., 1:2],
            torch.full_like(bowl[..., 2:3], 0.13),
        ],
        dim=-1,
    )
    current_td = policy_state["td"]
    cube = current_td["cube_pos"].clone()
    pinch_to_cube = current_td["pinch_pos"].clone() - cube
    position, quaternion = pose_at(drop_cube + pinch_to_cube, quat)
    # Action 10: Lower the grasped cube into the bowl.
    yield RobotMacroAction.reach_pose(
        position=position,
        quaternion=quaternion,
        gripper="closed",
        gripper_command=gripper_close_command,
        steps=100,
        settle_steps=40,
    )

    # Action 11: Open the gripper to release the cube.
    yield RobotMacroAction.open_gripper(steps=100, settle_steps=20)

    # Action 12: Wait with the gripper open so the cube can settle in the bowl.
    yield RobotMacroAction.wait(gripper="open", steps=240)

    current_td = policy_state["td"]
    retreat_xyz = torch.cat(
        [
            current_td["pinch_pos"][..., :2],
            torch.full_like(current_td["pinch_pos"][..., 2:3], 0.26),
        ],
        dim=-1,
    )
    position, quaternion = pose_at(retreat_xyz, quat)
    # Action 13: Retreat upward and away from the released cube.
    yield RobotMacroAction.reach_pose(
        position=position,
        quaternion=quaternion,
        gripper="open",
        steps=120,
        settle_steps=60,
    )

    # Action 14: Return the arm to the environment's reset joint configuration.
    yield RobotMacroAction.RESET

    # Action 15: Hold the reset pose while the final reward is measured.
    yield RobotMacroAction.wait(gripper="open", steps=800)


def action_mode(action) -> int | None:
    if not hasattr(action, "mode"):
        return None
    return int(action.mode.reshape(-1)[0].item())


def action_gripper(action) -> int:
    if not hasattr(action, "gripper"):
        return RobotMacroAction.GRIPPER_OPEN
    return int(action.gripper.reshape(-1)[0].item())


def run_scripted_rollout(td: TensorDictBase) -> TensorDictBase:
    policy_state["td"] = td
    policy_state["gripper_quat"] = td["pinch_quat"].clone()
    policy_state.pop("carry_start_cube", None)
    grasp_distance = torch.full_like(td["cube_pos"][..., :1], float("inf"))
    cube_motion_while_closed = torch.zeros_like(grasp_distance)
    cube_lift_while_closed = torch.zeros_like(grasp_distance)
    max_reward = torch.zeros_like(grasp_distance)
    last_reward = torch.zeros_like(grasp_distance)
    closed_reference_cube: torch.Tensor | None = None

    script_actions = gen_actions(td)
    while True:
        try:
            td["action"] = next(script_actions)
        except StopIteration:
            break

        action = td["action"]
        transition = env.step(td)
        reward = transition["next", "reward"]
        max_reward = torch.maximum(max_reward, reward.amax(dim=-2))
        last_reward = reward[..., -1, :]
        td = step_mdp(transition)
        policy_state["td"] = td

        mode = action_mode(action)
        gripper = action_gripper(action)
        if mode == int(RobotMacroActionMode.CLOSE_GRIPPER):
            grasp_distance = torch.minimum(
                grasp_distance, env.gripper_cube_distance(td)
            )
            closed_reference_cube = td["cube_pos"].clone()
        if (
            closed_reference_cube is not None
            and gripper == RobotMacroAction.GRIPPER_CLOSED
        ):
            cube_motion = (td["cube_pos"] - closed_reference_cube).norm(
                dim=-1, keepdim=True
            )
            cube_lift = td["cube_pos"][..., 2:3] - closed_reference_cube[..., 2:3]
            cube_motion_while_closed = torch.maximum(
                cube_motion_while_closed, cube_motion
            )
            cube_lift_while_closed = torch.maximum(cube_lift_while_closed, cube_lift)

    robot_home_error = (td["robot_qpos"] - home_qpos).norm(dim=-1, keepdim=True)

    assert grasp_distance.item() <= 0.025, grasp_distance.item()
    assert cube_motion_while_closed.item() >= 0.05, cube_motion_while_closed.item()
    assert cube_lift_while_closed.item() >= 0.08, cube_lift_while_closed.item()
    assert robot_home_error.item() <= 0.04, robot_home_error.item()
    assert max_reward.item() == 1.0, (
        max_reward.item(),
        td["cube_pos"],
        td["bowl_pos"],
    )
    assert last_reward.item() == 1.0, last_reward.item()
    assert td["success"].all()
    return td

Fixed scene video#

First, render the default task once. This is the simple cube-to-bowl scene used above to introduce the scripted policy.

td = run_scripted_rollout(td)

scripted_macro_animation = recorder.to_animation(
    title="Scripted RobotMacroAction rollout",
    interval=VIDEO_INTERVAL_MS,
    clear=True,
)

Randomized workspace video#

Next, reset the same environment onto new reachable layouts sampled from a square in the table plane. For readability, each rollout starts from a different pair of workspace regions: front-to-back, right-front to left-back, front-to-left-back, and right-front to center-back. The VideoRecorder keeps appending frames until we ask it to render, so the final video is one long clip containing all four randomized rollouts.

# In the CI fast path a single randomized rollout is enough to demonstrate
# the reset-onto-new-layout mechanics.
rollout_count = 1 if TUTORIAL_FAST else len(workspace_layouts)
for rollout_index in range(rollout_count):
    td = env.reset(random_scene_reset(rollout_index))
    td = run_scripted_rollout(td)

Rendered randomized rollouts#

The second animation shows the same hand-written policy on the randomized scenes.

randomized_macro_animation = recorder.to_animation(
    title="Randomized RobotMacroAction rollouts",
    interval=VIDEO_INTERVAL_MS,
    clear=True,
)

Reset and reuse the same environment#

The same transformed environment can be reset and reused for another scripted rollout.

reset_td = env.reset()
assert reset_td["robot_qpos"].shape == td["robot_qpos"].shape

Conclusion and further reading#

Macro actions are useful when a policy should choose what to do at a slower time scale while a scripted controller handles short-horizon actuation. They are not a replacement for learning in difficult contact-rich settings: they still depend on calibration, solver quality and reachable waypoints. Their practical value is that they provide a strong baseline, a source of demonstrations, and a structured action space for residual policies or RL fine tuning. Use gripper-only actions or wait when the arm must remain still; use an orientation mask and a Cartesian path when pose constraints must hold while the arm is moving.