# RoboMME

[RoboMME](https://robomme.github.io) is a memory-augmented manipulation benchmark built on ManiSkill (SAPIEN). It evaluates a robot's ability to retain and use information across an episode — counting, object permanence, reference, and imitation.

- **16 tasks** across 4 memory-skill suites
- **1,600 training demos** (100 per task, 50 val, 50 test)
- **Dataset**: [`lerobot/robomme`](https://huggingface.co/datasets/lerobot/robomme) — LeRobot v3.0, 768K frames at 10 fps
- **Simulator**: ManiSkill / SAPIEN, Panda arm, Linux only

![RoboMME benchmark tasks overview](https://cdn-thumbnails.huggingface.co/social-thumbnails/papers/2603.04639/gradient.png)

## Tasks

| Suite                             | Tasks                                                         |
| --------------------------------- | ------------------------------------------------------------- |
| **Counting** (temporal memory)    | BinFill, PickXtimes, SwingXtimes, StopCube                    |
| **Permanence** (spatial memory)   | VideoUnmask, VideoUnmaskSwap, ButtonUnmask, ButtonUnmaskSwap  |
| **Reference** (object memory)     | PickHighlight, VideoRepick, VideoPlaceButton, VideoPlaceOrder |
| **Imitation** (procedural memory) | MoveCube, InsertPeg, PatternLock, RouteStick                  |

## Installation

> RoboMME requires **Linux** (ManiSkill/SAPIEN uses Vulkan rendering). Docker is recommended to isolate dependency conflicts.

### Native (Linux)

```bash
pip install --override <(printf 'gymnasium==0.29.1\nnumpy==1.26.4\n') \
  -e '.[smolvla,av-dep]' \
  'robomme @ git+https://github.com/RoboMME/robomme_benchmark.git@main'
```

> **Dependency note**: `mani-skill` (pulled by `robomme`) pins `gymnasium==0.29.1` and `numpy<2.0.0`, which conflict with lerobot's base `numpy>=2.0.0`. That's why `robomme` is not a pyproject extra — use the override install above, or the Docker approach below to avoid conflicts entirely.

### Docker (recommended)

```bash
# Build base image first (from repo root)
docker build -f docker/Dockerfile.eval-base -t lerobot-eval-base .

# Build RoboMME eval image (applies gymnasium + numpy pin overrides)
docker build -f docker/Dockerfile.benchmark.robomme -t lerobot-robomme .
```

The `docker/Dockerfile.benchmark.robomme` image overrides `gymnasium==0.29.1` and `numpy==1.26.4` after lerobot's install. Both versions are runtime-safe for lerobot's actual API usage.

When evaluating a checkpoint saved on the host, run the container with the host UID. Checkpoint weights are intentionally private by default (`0600`), so the image's built-in user cannot otherwise read a bind-mounted `model.safetensors` file. Mount the Hugging Face cache at an accessible path as well:

```bash
docker run --gpus all --rm --ipc=host \
  --user "$(id -u):$(id -g)" \
  -e HF_HOME=/tmp/hf-cache \
  -v "$HOME/.cache/huggingface:/tmp/hf-cache:ro" \
  -v "$PWD/outputs:/results" \
  lerobot-robomme \
  lerobot-eval --policy.path=/results/<checkpoint>/pretrained_model # ...
```

## Running Evaluation

### Default (single task, single episode)

```bash
lerobot-eval \
    --policy.path=<your_policy_repo> \
    --env.type=robomme \
    --env.task=PickXtimes \
    --env.dataset_split=test \
    --env.task_ids=[0] \
    --eval.batch_size=1 \
    --eval.n_episodes=1
```

### Multi-task evaluation

Evaluate multiple tasks in one run by comma-separating task names. Use `task_ids` to select the dataset-backed episodes evaluated for every task. For the standard 50-episode test protocol, select IDs 0 through 49 and run each selected episode once.

```bash
lerobot-eval \
    --policy.path=<your_policy_repo> \
    --env.type=robomme \
    --env.task=PickXtimes,BinFill,StopCube,MoveCube,InsertPeg \
    --env.dataset_split=test \
    --env.task_ids=[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49] \
    --eval.batch_size=1 \
    --eval.n_episodes=1
```

### Key CLI options for `env.type=robomme`

| Option               | Default       | Description                                      |
| -------------------- | ------------- | ------------------------------------------------ |
| `env.task`           | `PickXtimes`  | Any of the 16 task names above (comma-separated) |
| `env.dataset_split`  | `test`        | `train`, `val`, or `test`                        |
| `env.action_space`   | `joint_angle` | `joint_angle` (8-D) or `ee_pose` (7-D)           |
| `env.episode_length` | `300`         | Max steps per episode                            |
| `env.task_ids`       | `null`        | Dataset episode indices (null = `[0]`)           |

`eval.n_episodes` repeats each selected `task_id`; it does not advance to the next dataset episode. Keep it at `1` when enumerating the 50 distinct test episodes with `env.task_ids`.

## Dataset

The dataset [`lerobot/robomme`](https://huggingface.co/datasets/lerobot/robomme) is in **LeRobot v3.0 format** and can be loaded directly:

```python
from lerobot.datasets.lerobot_dataset import LeRobotDataset

dataset = LeRobotDataset("lerobot/robomme")
```

### Dataset features

| Feature            | Shape         | Description                      |
| ------------------ | ------------- | -------------------------------- |
| `image`            | (256, 256, 3) | Front camera RGB                 |
| `wrist_image`      | (256, 256, 3) | Wrist camera RGB                 |
| `actions`          | (8,)          | Joint angles + gripper           |
| `state`            | (8,)          | Joint positions + gripper state  |
| `simple_subgoal`   | str           | High-level language annotation   |
| `grounded_subgoal` | str           | Grounded language annotation     |
| `exec_start_idx`   | scalar        | First action-execution frame     |
| `is_demo`          | bool          | Video-demonstration prefix frame |
| `episode_index`    | int           | Episode ID                       |
| `frame_index`      | int           | Frame within episode             |

### Feature key alignment (training)

The env wrapper exposes `pixels/image` and `pixels/wrist_image` as observation keys. The `features_map` in `RoboMMEEnv` maps these to `observation.images.image` and `observation.images.wrist_image` for the policy. State is exposed as `agent_pos` and maps to `observation.state`.

The published dataset uses the raw keys shown above. Policies that expect canonical LeRobot keys should map them at training time. For example, the pretrained SmolVLA checkpoint expects three cameras, while RoboMME provides two:

```bash
uv run lerobot-train \
  --policy.path=lerobot/smolvla_base \
  --policy.empty_cameras=1 \
  --dataset.repo_id=lerobot/robomme \
  --dataset.training_target_start_feature=exec_start_idx \
  '--rename_map={"image":"observation.images.camera1","wrist_image":"observation.images.camera2","state":"observation.state","actions":"action"}'
```

RoboMME episodes begin with video-demonstration/context frames that should remain available to a
memory policy but should not be sampled as action-training targets. Setting
`dataset.training_target_start_feature=exec_start_idx` starts target sampling at each episode's
execution boundary while preserving earlier frames for temporal observation deltas. In the published
training split this keeps 476,857 execution targets out of 768,897 total frames. One execution-target
epoch therefore takes `ceil(476857 / effective_batch_size)` optimizer steps.

For a sample-matched SmolVLA visual-memory ablation, use
`examples/robomme/smolvla_visual_memory_ablation.sh`. `TARGET_SAMPLES` counts examples across all
GPUs, and `NUM_PROCESSES` is included when the script converts that target into optimizer steps. For
example, the following reproduces 5.12 million example exposures (the exposure of RoboMME's
80,000-step, global-batch-64 memory-policy recipe) with four GPUs, two accumulated microbatches,
and an effective global batch of 64:

```bash
TARGET_SAMPLES=5120000 \
NUM_PROCESSES=4 \
BATCH_SIZE=8 \
GRADIENT_ACCUMULATION_STEPS=2 \
VARIANT=visual-memory \
RUN_EVAL=false \
bash examples/robomme/smolvla_visual_memory_ablation.sh
```

This becomes 80,000 optimizer steps, or about 10.74 execution-target epochs. Run the baseline with
the same `TARGET_SAMPLES`, effective batch size, seed, and scheduler settings to isolate the visual
memory change. RoboMME's released baseline uses a different global batch (128), so its nominal
80,000-step recipe is not sample-matched to its global-batch-64 memory recipe.

At evaluation time the environment wrapper has already converted observations to canonical keys. Only the two camera suffixes need to be aligned with the checkpoint:

```bash
'--rename_map={"observation.images.image":"observation.images.camera1","observation.images.wrist_image":"observation.images.camera2"}'
```

## Action Spaces

| Type          | Dim | Description                                               |
| ------------- | --- | --------------------------------------------------------- |
| `joint_angle` | 8   | 7 joint angles + 1 gripper (−1 closed, +1 open, absolute) |
| `ee_pose`     | 7   | xyz + roll/pitch/yaw + gripper                            |

Set via `--env.action_space=joint_angle` (default) or `--env.action_space=ee_pose`.

## Platform Notes

- **Linux only**: ManiSkill requires SAPIEN/Vulkan. macOS and Windows are not supported.
- **GPU recommended**: Rendering is CPU-capable but slow; CUDA + Vulkan gives full speed.
- **gymnasium / numpy conflict**: See installation note above. Docker image handles this automatically.
- **ManiSkill fork**: `robomme` depends on a specific ManiSkill fork (`YinpeiDai/ManiSkill`), pulled in automatically via the `robomme` package.