GitHub - Stanford-TML/robot_keyframe_kit: Create motion for any robot through editing keyframe (original) (raw)

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:

Detect that no floor plane exists
Check for a scene.xml in the same directory and suggest using it
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:

Command-Line Options

keyframe-editor [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