robot-keyframe-kit

A generalizable MuJoCo keyframe editor for creating and editing robot motion sequences. Works with any MuJoCo-compatible robot model.

🎬 Video Tutorial

▶️ Click to watch the full tutorial on YouTube

Installation

pip install robot-keyframe-kit

Development Setup (Conda)

For contributors, use an editable install so local code changes are reflected immediately.

git clone https://github.com/Stanford-TML/robot_keyframe_kit.git
cd robot_keyframe_kit

conda create -n robot_keyframe_kit python=3.10 -y
conda activate robot_keyframe_kit

python -m pip install -e ".[dev]"

# Optional: download MuJoCo Menagerie assets for local testing/examples
bash scripts/setup_assets.sh

Verify the editable install and CLI entrypoint:

python -c "import robot_keyframe_kit; print(robot_keyframe_kit.__file__)"
keyframe-editor --help

Optional example (after running bash scripts/setup_assets.sh):

keyframe-editor assets/mujoco_menagerie/toddlerbot_2xc/scene.xml --name toddlerbot_2xc

Quick Start

Using Scene XML (Recommended)

For best results, use a scene.xml file that includes your robot model along with proper floor setup.

keyframe-editor /path/to/scene.xml --name my_robot

Most robot models from mujoco_menagerie include a scene.xml file. For example:

keyframe-editor /path/to/mujoco_menagerie/unitree_g1/scene.xml --name g1

Using Robot-Only XML

If you only have a robot XML file (without scene setup), the editor will automatically:

1. Detect that no floor plane exists
2. Check for a scene.xml in the same directory and suggest using it
3. Auto-generate a scene wrapper with floor plane for physics collision

keyframe-editor /path/to/robot.xml --name my_robot

Note: For better visualization and physics, create or use a proper scene.xml.

Creating a Scene XML

If your robot model doesn't have a scene.xml, you can create one:

<mujoco model="my_robot_scene">
  <!-- Include your robot model -->
  <include file="robot.xml"/>

  <!-- Add floor and lighting -->
  <asset>
    <texture type="2d" name="groundplane" builtin="checker" 
             rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3"
             markrgb="0.8 0.8 0.8" width="300" height="300"/>
    <material name="groundplane" texture="groundplane" 
              texuniform="true" texrepeat="5 5" reflectance="0.2"/>
  </asset>

  <worldbody>
    <light pos="0 0 3.5" dir="0 0 -1" directional="true"/>
    <geom name="floor" type="plane" size="10 10 0.05" material="groundplane"/>
  </worldbody>
</mujoco>

Command-Line Options

keyframe-editor <xml_path> [OPTIONS]

Required Arguments

• xml_path: Path to the MuJoCo XML file (scene.xml recommended, or robot.xml)

Optional Arguments

• --name <name>: Name for this project (used in save filenames). Default: robot
• --root-body <body_name>: Name of the root body for ground alignment. Auto-detected if not specified.
• --config <path>: Path to YAML configuration file. Overrides auto-detection.
• --generate-config <path>: Generate a YAML configuration file from the model and exit.
• --data <path>: Path to load existing keyframe data from.
• --save-dir <dir>: Directory to save keyframe data. Default: keyframes
• --no-auto-floor: Disable automatic floor injection for robot-only XML files.

Configuration Files

You can create a YAML configuration file to customize robot-specific settings:

name: my_robot
root_body: base_link
end_effector_sites:
  - left_foot
  - right_foot
mirror_pairs:
  left_hip_pitch: right_hip_pitch
  left_knee: right_knee
mirror_signs:
  left_hip_pitch: -1
  left_knee: -1
dt: 0.02
save_dir: keyframes
scene:
  auto_inject_floor: true
  show_floor: true

Generate a default config from your model:

keyframe-editor robot.xml --generate-config config.yaml

Features

• Visual-Centric Joint Control: Control motion joints intuitively
• Automatic Mechanism Detection: Handles differential drives, gear couplings, and parallel linkages
• Mirror Mode: Automatically mirrors joint movements with correct sign conventions
• Physics Simulation: Test keyframes and trajectories with full MuJoCo physics
• Ground Placement: Automatically places robot on ground based on lowest collision geometry
• End-Effector Tracking: Auto-detects and tracks end-effector sites/bodies
• Mink IK Solver: QP-based IK for end-effector target solving

Python API

You can also use the editor programmatically:

from robot_keyframe_kit import ViserKeyframeEditor, EditorConfig

# Load config from file
config = EditorConfig.from_yaml("config.yaml")

# Or create config programmatically
config = EditorConfig(
    name="my_robot",
    root_body="base_link",
    auto_inject_floor=True,
    show_floor=True,
)

# Create editor
editor = ViserKeyframeEditor(
    "scene.xml",
    config=config,
)

# Editor runs until interrupted

Save File Format

Motion data is saved as LZ4-compressed pickle files (.lz4) using joblib. Files are saved to {save_dir}/{name}/{motion_name}.lz4.

Loading Save Files

import joblib

data = joblib.load("keyframes/my_robot/walk.lz4")

File Structure

Keyframe Dictionary Structure

Each keyframe in the keyframes list contains:

Example Usage

import joblib
import numpy as np

# Load motion data
data = joblib.load("keyframes/toddlerbot/wave.lz4")

# Get keyframes
for kf in data["keyframes"]:
    print(f"Keyframe: {kf['name']}, motor_pos shape: {kf['motor_pos'].shape}")

# Get trajectory
times = data["time"]  # (T,)
qpos = data["qpos"]   # (T, nq)
print(f"Trajectory: {len(times)} frames, {times[-1]:.2f}s duration")

# Get timed sequence for playback
for keyframe_name, duration in data["timed_sequence"]:
    print(f"  {keyframe_name}: {duration}s")

Troubleshooting

Robot Falls Through Floor

• Check: Does your XML have a floor plane? Use scene.xml if available.
• Solution: The editor auto-injects a floor, but collision filtering may need adjustment.
• Best Fix: Use a proper scene.xml with correct collision settings.

Slow Physics Simulation

• Check: Model timestep settings (model.opt.timestep)
• Solution: Editor uses n_frames=20 substeps per control step for stability.

Wrong Joint Selection

• Check: Are you seeing motor joints instead of motion joints?
• Solution: The editor auto-detects differential drives and gear mechanisms. Check your config for manual overrides.

Citation

If you find this tool useful for your research, please consider citing:

@misc{yang2026locomotion,
  title = {Locomotion {{Beyond Feet}}},
  author = {Yang, Tae Hoon and Shi, Haochen and Hu, Jiacheng and Zhang, Zhicong and Jiang, Daniel and Wang, Weizhuo and He, Yao and Wu, Zhen and Chen, Yuming and Hou, Yifan and Kennedy, Monroe and Song, Shuran and Liu, C. Karen},
  year = 2026,
  month = jan,
  number = {arXiv:2601.03607},
  eprint = {2601.03607},
  primaryclass = {cs},
  publisher = {arXiv},
  doi = {10.48550/arXiv.2601.03607},
  urldate = {2026-01-08},
  archiveprefix = {arXiv},
  keywords = {Computer Science - Robotics}
}

@article{shi2025toddlerbot,
  title={ToddlerBot: Open-Source ML-Compatible Humanoid Platform for Loco-Manipulation},
  author={Shi, Haochen and Wang, Weizhuo and Song, Shuran and Liu, C. Karen},
  journal={arXiv preprint arXiv:2502.00893},
  year={2025}
}

License

MIT License