🚧 Early Prototype — Community Feedback Welcome!

This project is currently in an early prototype stage.

Features, architecture, and documentation are actively evolving, and breaking changes are likely as we iterate.

We’re building this openly with the community, so feedback, ideas, and contributions are highly encouraged! If you’d like to help shape the direction of the project:

Open an issue to share suggestions or report bugs Start a discussion about improvements Submit a pull request with enhancements Thank you for helping us improve this project!

MyoSuite to mjlab Integration

Integration package for using MyoSuite environments with mjlab's training infrastructure.

Features

• ✅ Automatic Registration: All MyoSuite environments are automatically registered with mjlab
• ✅ MJX/Warp Support: Compatible with both standard MyoSuite and mjx/warp GPU-accelerated versions
• ✅ Native mjlab Integration: Uses mjlab's native task registration when available
• ✅ Backward Compatible: Falls back to gymnasium registry if mjlab native registration unavailable
• ✅ Full Test Coverage: Comprehensive unit tests for all functionality

Installation

# Install mjlab-myosuite
pip install -e .
pip install "myosuite @ git+https://github.com/MyoHub/myosuite.git@mjx"

# Or with uv (faster)
uv venv
uv pip install -e .
uv pip install "myosuite @ git+https://github.com/MyoHub/myosuite.git@mjx"

Quick Start

1. Basic Usage (gym based)

import gymnasium as gym
import mjlab_myosuite  # Auto-registers all MyoSuite environments

# Create a MyoSuite environment wrapped for mjlab
env = gym.make("Mjlab-MyoSuite-myoElbowPose1D6MRandom-v0")
obs, info = env.reset()
action = env.action_space.sample()
obs, rewards, dones, extras = env.step(action)
env.close()

2. Training with mjlab

Use the mjlab training scripts:

# Train a policy
uv run train Mjlab-MyoSuite-myoElbowPose1D6MRandom-v0 \
    --agent.max-iterations 200 \
    --agent.num-steps-per-env 512

# Play with trained policy
uv run play Mjlab-MyoSuite-myoElbowPose1D6MRandom-v0 \
    --checkpoint_file logs/rsl_rl/myosuite/.../model_199.pt

3. Custom Task Registration

For registering specific tasks with custom configurations, see:

• examples/example_task_registration.py - Complete example following mjlab's tutorial pattern

Architecture

The integration follows mjlab's native task registration pattern from the create_new_task tutorial:

┌─────────────────────────────────────┐
│   mjlab Training Pipeline           │
│   (PPO, WandB, etc.)                │
└──────────────┬──────────────────────┘
               │
┌──────────────▼──────────────────────┐
│   MyoSuite Wrapper                  │
│   - Adapts Gym API to mjlab         │
│   - Handles batched operations      │
└──────────────┬──────────────────────┘
               │
┌──────────────▼──────────────────────┐
│   MyoSuite (Standard or MJX/Warp)   │
│   - Musculoskeletal models          │
│   - Task-specific rewards           │
└─────────────────────────────────────┘

Supported MyoSuite Versions

• Standard MyoSuite: CPU-based MuJoCo simulation
• MJX/Warp MyoSuite: GPU-accelerated from the mjx branch

The wrapper automatically detects and supports both versions.

Configuration

Environment Configuration

from mjlab_myosuite.config import MyoSuiteEnvCfg

cfg = MyoSuiteEnvCfg()
cfg.num_envs = 4096  # For training
cfg.device = "cuda:0"  # Use GPU for mjx/warp versions

RL Configuration

from mjlab_myosuite.config import get_default_myosuite_rl_cfg

rl_cfg = get_default_myosuite_rl_cfg()
rl_cfg.max_iterations = 2000
rl_cfg.algorithm.learning_rate = 3e-4

Known Issues

TBD

ONNX Model Export

MyoSuite environments support ONNX model export for deployment and inference. The MyoSuiteOnPolicyRunner automatically exports ONNX models when using wandb logging:

from mjlab_myosuite.rl.runner import MyoSuiteOnPolicyRunner

# During training, ONNX models are automatically exported
runner = MyoSuiteOnPolicyRunner(env, agent_cfg, log_dir, device)
runner.learn(num_learning_iterations=1000)
# ONNX model is saved alongside the PyTorch checkpoint

The exported ONNX model includes:

• Policy network (actor) with optional observation normalizer
• MyoSuite-specific metadata (action dimensions, observation dimensions, etc.)
• Compatibility with ManagerBasedRlEnv structure

Tracking Tasks (not yet fully implemented)

MyoSuite tracking tasks follow the same structure as mjlab's tracking tasks, allowing you to train policies to track reference motions. The tracking functionality is implemented in src/mjlab_myosuite/tasks/tracking/ following the mjlab tracking structure.

Tracking Configuration

from mjlab_myosuite.tasks.tracking.tracking_env_cfg import MyoSuiteTrackingEnvCfg

# Create tracking configuration
cfg = MyoSuiteTrackingEnvCfg()
cfg.num_envs = 4096
cfg.device = "cuda:0"
cfg.commands.motion.motion_file = "path/to/motion.npz"

Training Tracking Tasks

# Train with motion file from wandb artifact
uv run train Mjlab-MyoSuite-Tracking-myoElbowPose1D6MRandom-v0 \
    --motion-file examples/elbow_sinusoidal_motion.npz     \
    --agent.max-iterations 10000

Playing Tracking Tasks

# Play with motion file
uv run play Mjlab-MyoSuite-Tracking-myoElbowPose1D6MRandom-v0 \
    --checkpoint_file logs/rsl_rl/myosuite/.../model_2000.pt \
    --motion-file path/to/motion.npz

The tracking runner (MyoSuiteMotionTrackingOnPolicyRunner) extends the base MyoSuite runner and provides support for motion tracking, including wandb artifact integration for motion files.

Viser Playback Utility

The playback_with_viser utility provides a convenient way to visualize policy execution using the Viser web-based viewer:

from scripts.play import playback_with_viser
import gymnasium as gym

# Create environment and policy
env = gym.make("Mjlab-MyoSuite-myoElbowPose1D6MRandom-v0")
policy = load_policy("path/to/checkpoint.pt")

# Playback with Viser
playback_with_viser(env, policy, verbose=True)

You can also use it from the command line:

# Use Viser viewer explicitly
play Mjlab-MyoSuite-myoElbowPose1D6MRandom-v0 \
    --viewer viser \
    --checkpoint_file logs/rsl_rl/myosuite/.../model_2000.pt

# Specify Viser server port
play Mjlab-MyoSuite-myoElbowPose1D6MRandom-v0 \
    --viewer viser \
    --viser-port 8080 \
    --checkpoint_file logs/rsl_rl/myosuite/.../model_2000.pt

Testing

# Run all tests
pytest tests/

# Run specific test
pytest tests/test_myosuite_integration.py::test_wrapper_creation_direct

# Run GPU acceleration tests (requires CUDA)
pytest tests/test_gpu_acceleration.py -v

# Run ONNX export tests (requires ONNX)
pytest tests/test_onnx_export.py -v

Development

Run tests:

make test          # Run all tests
make test-fast     # Skip slow integration tests
uv run --no-default-groups --group cu128 --group dev pyright
uv run --no-default-groups --group cu128 --group dev pytest

Format code:

uvx pre-commit install
make format

Documentation

• mjlab Tutorial - Official mjlab task creation tutorial
• MyoSuite Documentation - MyoSuite documentation