|
| 1 | +# Multi-Agent (Bimanual) Datasets |
| 2 | + |
| 3 | +RoboVerse's trajectory format is multi-agent native. A dataset file stores one |
| 4 | +entry **per agent, keyed by robot name** — the same on-disk layout single-agent |
| 5 | +datasets already use. A single-agent file is therefore just the one-key special |
| 6 | +case, so existing datasets keep working unchanged. |
| 7 | + |
| 8 | +This is what makes bimanual workflows (two independent arms acting |
| 9 | +simultaneously, e.g. ManiSkill's `TwoRobotStackCube-v1` style tasks) expressible |
| 10 | +without inventing a parallel format. |
| 11 | + |
| 12 | +## On-disk format |
| 13 | + |
| 14 | +A `*_v2.pkl` file is a dict keyed by robot name. Each agent maps to a list of |
| 15 | +demos; each demo carries `init_state`, `actions`, and optional `states`: |
| 16 | + |
| 17 | +```python |
| 18 | +{ |
| 19 | + "franka_left": [{"init_state": {...}, "actions": [...], "states": None}, ...], |
| 20 | + "franka_right": [{"init_state": {...}, "actions": [...], "states": None}, ...], |
| 21 | + "metadata": {"num_agents": 2, "agents": ["franka_left", "franka_right"]}, |
| 22 | +} |
| 23 | +``` |
| 24 | + |
| 25 | +Each agent's `init_state` lists that agent's robot entry plus any **shared |
| 26 | +objects** (the cube both arms coordinate around). Per-agent actions are |
| 27 | +namespaced as `{"dof_pos_target": {...}}`. |
| 28 | + |
| 29 | +## Loading with `get_traj` |
| 30 | + |
| 31 | +The canonical loader `metasim.utils.demo_util.get_traj` takes either a single |
| 32 | +robot (single-agent, unchanged) or a **list of robots** (multi-agent): |
| 33 | + |
| 34 | +```python |
| 35 | +from metasim.utils.demo_util import get_traj |
| 36 | + |
| 37 | +robots = [franka.replace(name=n) for n in ["franka_left", "franka_right"]] |
| 38 | +init_states, all_actions, all_states = get_traj("bimanual_handover_v2.pkl", robots) |
| 39 | +``` |
| 40 | + |
| 41 | +Passing the list returns the **same three-tuple shape** as the single-agent |
| 42 | +path, with every agent merged into each per-step dict: |
| 43 | + |
| 44 | +- `init_states[d]["robots"]` holds every arm; `init_states[d]["objects"]` holds |
| 45 | + the shared objects once. |
| 46 | +- `all_actions[d][t]` is `{robot_name: {"dof_pos_target": ...}}` for **all** |
| 47 | + agents at step `t` — exactly what `handler.set_dof_targets([...])` consumes. |
| 48 | +- `all_states[d][t]` unions each agent's `robots`/`objects` (or is `None` for |
| 49 | + action-only demos). |
| 50 | + |
| 51 | +Because the shape is identical, the same replay / collection code paths drive |
| 52 | +one arm or many. Multi-agent loading requires the v3 namespaced format |
| 53 | +(`v2_as_v3=True`, the default); `v2_as_v3=False` with a robot list raises, since |
| 54 | +namespacing is what keeps each agent's actions indexed by name. |
| 55 | + |
| 56 | +## Runnable examples |
| 57 | + |
| 58 | +`get_started/8_multiagent_dataset.py` builds a coordinated two-Franka handover |
| 59 | +trajectory, saves it as a real `*_v2.pkl`, loads it back through `get_traj`, and |
| 60 | +replays both arms simultaneously to video: |
| 61 | + |
| 62 | +```bash |
| 63 | +MUJOCO_GL=egl python get_started/8_multiagent_dataset.py --sim mujoco |
| 64 | +``` |
| 65 | + |
| 66 | +The same trajectory is also exposed as a registered task, so it replays through |
| 67 | +the **canonical pipeline** (`scripts/advanced/replay_demo.py`) — which now passes |
| 68 | +the full robot list to `get_traj` whenever a task declares more than one robot: |
| 69 | + |
| 70 | +```bash |
| 71 | +MUJOCO_GL=egl python scripts/advanced/replay_demo.py \ |
| 72 | + --task bimanual.franka_handover --sim mujoco --headless |
| 73 | +``` |
| 74 | + |
| 75 | +`get_started/9_maniskill_two_robot_stack_cube.py` does the same round trip with |
| 76 | +**real ManiSkill data**: it fetches the official `TwoRobotStackCube-v1` |
| 77 | +demonstrations, converts one episode into the name-keyed `*_v2` format, loads |
| 78 | +both Panda arms through `get_traj`, and replays the recorded states on MuJoCo: |
| 79 | + |
| 80 | +```bash |
| 81 | +MUJOCO_GL=egl python get_started/9_maniskill_two_robot_stack_cube.py --sim mujoco |
| 82 | +``` |
| 83 | + |
| 84 | +The ManiSkill `.h5` stores one articulation per agent |
| 85 | +(`panda_wristcam-agent-0` / `-agent-1`) plus the shared cubes; converting it is |
| 86 | +just a regrouping into one keyed entry per agent. Replay uses the recorded |
| 87 | +**states** (kinematic playback) rather than open-loop action targets: the demos |
| 88 | +were collected under SAPIEN's `pd_joint_delta_pos` controller, and closed-loop |
| 89 | +contact dynamics do not transfer across simulators, so state replay is the |
| 90 | +faithful cross-sim view of the dataset. |
| 91 | + |
| 92 | +## Single-embodiment bimanual vs. two agents |
| 93 | + |
| 94 | +Two distinct cases share this format: |
| 95 | + |
| 96 | +- **Single-embodiment bimanual** (one URDF with two arms, e.g. ALOHA / RoboTwin |
| 97 | + AgileX) — one robot entry whose action dict spans all joints. See the |
| 98 | + [RoboTwin Integration](../integrations/robotwin.md). |
| 99 | +- **Two independent agents** (two separate robot entities) — the case above, |
| 100 | + one keyed entry per agent. |
| 101 | + |
| 102 | +The single-embodiment bimanual case is demonstrated by |
| 103 | +`get_started/10_robotwin_aloha_replay.py`, but note: |
| 104 | + |
| 105 | +```{warning} |
| 106 | +`get_started/10_robotwin_aloha_replay.py` is **experimental** and **not |
| 107 | +out-of-the-box** (unlike examples 8 and 9, which run from a clean MuJoCo |
| 108 | +install). It needs a local RoboTwin clone, its ~3.74 GB asset pack, a separate |
| 109 | +`robotwin` conda env, and a curobo build for the local GPU arch to collect a |
| 110 | +bridge pickle first. The manipulated object is rendered as a primitive-cube |
| 111 | +proxy (not the real mesh), and only joint *motion* — not task *success* — has |
| 112 | +been confirmed. Treat it as a data-bridge demo, not a benchmark. |
| 113 | +``` |
0 commit comments