robot-keyframe-kit 0.3.0

A generalizable Viser-based keyframe editor for any MuJoCo robot

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

Key	Type	Description
keyframes	List[dict]	List of keyframe dictionaries (see below)
timed_sequence	List[Tuple[str, float]]	Sequence of (keyframe_name, duration_sec) pairs
time	ndarray (T,)	Timestamps for each trajectory frame
qpos	ndarray (T, nq)	Full MuJoCo qpos at each frame
motor_vel	ndarray (T, n_motors)	Motor/actuator joint velocities (rad/s)
joint_vel	ndarray (T, n_joints)	UI-visible joint velocities (rad/s)
action	ndarray (T, nu) or None	Motor commands (if action trajectory was played)
body_pos	ndarray (T, nbody, 3)	Body positions (world or relative frame)
body_quat	ndarray (T, nbody, 4)	Body orientations as quaternions (w, x, y, z)
body_lin_vel	ndarray (T, nbody, 3)	Body linear velocities
body_ang_vel	ndarray (T, nbody, 3)	Body angular velocities
site_pos	ndarray (T, n_sites, 3)	End-effector site positions
site_quat	ndarray (T, n_sites, 4)	End-effector site orientations
is_robot_relative_frame	bool	Whether poses are in robot-relative frame

Keyframe Dictionary Structure

Each keyframe in the keyframes list contains:

Key	Type	Description
name	str	Human-readable keyframe name
motor_pos	ndarray (nu,)	Motor/actuator positions
joint_pos	ndarray (nj,) or None	Joint positions (may differ from motor_pos with transmissions)
qpos	ndarray (nq,) or None	Full MuJoCo qpos including base pose

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