MuscleMimic

repository·main·Indexed 42 days ago

https://github.com/amathislab/musclemimic

A JAX-based motion imitation learning benchmark for biomechanically accurate, muscle-actuated models. It enables full-body locomotion and manipulation training using GPU-parallelized environments via MuJoCo/MJWarp, featuring models such as MyoBimanualArm and MyoFullBody.

Tokens
6K
Snippets
19
Records
29
Agent score
56%

What's inside MuscleMimic

  1. How the C3D to SMPL pipeline works

    main

    The C3D to SMPL pipeline is inspired by MoSh++. It follows a two-stage optimization process:

    1. Stage I: Optimizes shared betas (shape), per-reference-frame rigid alignment, and markers_latent (latent marker positions). This stage uses a dense Powell dogleg optimizer (available in both Torch and JAX).
    2. Stage II: Keeps shape and latent markers fixed, then fits pose and translation frame by frame.

    Key Pipeline Steps:

    • Canonicalize marker labels to MoSh++ names.
    • Attach markers to SMPL-X/H surface locations using MoSh++ vertex IDs.
    • Initialize markers_latent from mesh surface normals.
    • Export results to an AMASS/SMPL-H-compatible format for GMR.
  2. Manage the converted C3D cache

    main

    The C3D pipeline outputs an extended MuscleMimic Trajectory .npz file. These are stored in a cache directory.

    • Default location: ~/.musclemimic/caches/C3D/<model>/<method>/<dataset-name>.npz.
    • Custom location: Set via musclemimic-set-conv-c3d-path --path /path/to/dir or the MUSCLEMIMIC_CONVERTED_C3D_PATH environment variable.
    • Dataset Naming: Use --c3d-dataset-name to specify a relative name under the cache root. If omitted, a name is generated based on the file path and a hash of settings.
    • Recomputing: If you need to re-run the pipeline and ignore existing cache, use the --clear-c3d-cache flag.
  3. Install MuscleMimic

    main

    MuscleMimic requires Linux with an NVIDIA GPU for training. Inference and evaluation are supported on both Linux and macOS.

    To install, use uv to manage dependencies and include the cuda extra:

    curl -LsSf https://astral.sh/uv/install.sh | sh
    git clone https://github.com/amathislab/musclemimic
    cd musclemimic
    uv sync --extra cuda
  4. Install the MuscleMimic Web Viewer

    main

    The installation command depends on the level of functionality required. Use uv sync with the appropriate extras:

    • Raw C3D viewer only: Installs minimal dependencies for viewing markers.
    • C3D fitting and default GMR retargeting: Installs dependencies for fitting markers to SMPL surfaces and retargeting via GMR.
    • CUDA-accelerated JAX Stage I: Adds GPU acceleration for the Stage-I solver.
    # Raw C3D viewer only
    uv sync --extra c3d
    
    # C3D fitting and default GMR retargeting
    uv sync --extra c3d --extra smpl --extra gmr
    
    # CUDA-accelerated default JAX Stage I
    uv sync --extra c3d --extra smpl --extra gmr --extra cuda
  5. Set up demo motions for MuscleMimic

    main

    To run a demo without manually downloading AMASS datasets, you can use pre-retargeted demo motions. Note that the demo dataset is gated on Hugging Face, so you must first authenticate with Hugging Face.

    1. Authenticate with Hugging Face: uv run hf auth login

    2. Setup the demo cache for your target model:

      • For MyoBimanualArm: uv run python -c "from musclemimic.utils.demo_cache import setup_demo_for_bimanual; setup_demo_for_bimanual()"
      • For MyoFullBody: uv run python -c "from musclemimic.utils.demo_cache import setup_demo_for_myo_fullbody; setup_demo_for_myo_fullbody()"
    uv run hf auth login
    uv run python -c "from musclemimic.utils.demo_cache import setup_demo_for_bimanual; setup_demo_for_bimanual()"
    uv run python -c "from musclemimic.utils.demo_cache import setup_demo_for_myo_fullbody; setup_demo_for_myo_fullbody()"
  6. Configure SMPL/SMPL-X Model Paths

    main

    SMPL/SMPL-X model files are not included in this repository and must be downloaded from official sources. You must provide these paths for C3D fitting or retargeting.

    For C3D Fitting

    Use --c3d-model-path or configure the environment variable MUSCLEMIMIC_C3D_MODEL_PATH using the CLI tool. The path should point to a directory containing the model files (e.g., male/model.pkl, SMPLX_MALE.pkl, or SMPLX_MALE.npz).

    For Direct SMPL Retargeting

    If using --retargeting-method smpl, you must also provide a retargeting-compatible SMPL/SMPL-H model path via --retarget-smpl-model-path or the MUSCLEMIMIC_SMPL_MODEL_PATH environment variable.

    # Set C3D fitting model path
    musclemimic-set-c3d-model-path --path /path/to/c3d-fitting/smplx-or-smplh
    
    # Set direct SMPL retargeting model path
    musclemimic-set-smpl-model-path --path /path/to/retarget/smplh/models
  7. Submit MuscleMimic jobs with Run:AI

    main

    To run MuscleMimic experiments on a Run:AI cluster, use the runai submit command. You must provide a container image containing the MuscleMimic environment, specify your UID/GID, and configure environment variables for uv (the package manager) to ensure dependencies are managed correctly within the persistent volume.

    Key configuration requirements:

    • Image: Use your pre-built MuscleMimic container image.
    • GPU/Node Pools: Specify the required number of GPUs and the target node pool (e.g., h100).
    • Storage: Use --existing-pvc to mount a persistent volume, typically at /users.
    • Environment Variables: Set UV_PROJECT_ENVIRONMENT, UV_CACHE_DIR, UV_PYTHON_INSTALL_DIR, and HOME to point to paths within your mounted PVC to ensure persistence across job restarts.
    • Execution: The command should cd into your project directory, run uv sync with the necessary extras (smpl, gmr, cuda), and then execute your experiment script using uv run.
    runai submit \
      --name mm-$(date +%m%d-%H%M%S) \
      --image <your-musclemimic-image> \
      --run-as-uid <your_epfl_uid> \
      --run-as-gid <your_epfl_gid> \
      --gpu 1 --node-pools h100 \
      --existing-pvc claimname=<your_claim_name>,path=/users \
      --environment UV_PROJECT_ENVIRONMENT=/tmp/venv \
      --environment UV_CACHE_DIR=/users/<your_username>/.cache/uv \
      --environment UV_PYTHON_INSTALL_DIR=/users/<your_username>/.uv-pythons \
      --environment HOME=/users/<your_username> \
      --environment WANDB_MODE=disabled \
      --backoff-limit 0 \
      --command -- /bin/bash -c "cd /users/<your_username>/musclemimic; uv sync --extra smpl --extra gmr --extra cuda; uv run fullbody/experiment.py --config-name=<some_config>"
  8. Configure MoSh++ Assets for Fitting

    main

    To achieve MoSh++-parity fitting, you can provide additional assets like body-pose priors and head marker correlations. You can configure a single directory containing these files using musclemimic-set-moshpp-assets-path (sets MUSCLEMIMIC_MOSHPP_ASSETS_PATH), or set individual paths for each asset:

    • MUSCLEMIMIC_MOSHPP_POSE_HAND_PRIOR_PATH: Path to pose_hand_prior.npz.
    • MUSCLEMIMIC_MOSHPP_POSE_BODY_PRIOR_PATH: Path to pose_body_prior.pkl.
    • MUSCLEMIMIC_MOSHPP_HEAD_MARKER_CORR_PATH: Path to ssm_head_marker_corr.npz.
    # Set a directory containing all MoSh++ assets
    musclemimic-set-moshpp-assets-path --path /path/to/moshpp/assets
  9. Use the musclemimic-retarget-visualize CLI

    main

    The musclemimic-retarget-visualize command is used to visualize retargeted motion or generate videos from C3D files on musculoskeletal models.

    Input Modes You must provide exactly one of the following input methods:

    1. --motion: Specify one or more motion names (e.g., 'KIT/3/tennis_forehand_right04_poses'). This can be passed multiple times.
    2. --dataset-group: Specify a dataset group from loco_mujoco.smpl.const.
    3. --c3d-file: Provide a path to a .c3d file. This mode automatically fits SMPL, retargets, and visualizes the motion.

    Common Usage Examples

    Visualize specific AMASS motions and record a video:

    uv run musclemimic-retarget-visualize --motion "KIT/6/WalkInCounterClockwiseCircle06_1_poses" --record

    Visualize a C3D file with specific extras (C3D, SMPL, and GMR support):

    uv run --extra c3d --extra smpl --extra gmr musclemimic-retarget-visualize --c3d-file path/to/file.c3d --c3d-model-path /path/to/smplx
    uv run musclemimic-retarget-visualize --motion "KIT/6/WalkInCounterClockwiseCircle06_1_poses" --record
  10. Fit C3D markers and retarget via GMR

    main

    To fit C3D markers to a SMPL surface and retarget the motion using the default GMR method, use the musclemimic.web_viewer.run module. You must provide the --c3d-file, --c3d-model-path, and a --c3d-dataset-name.

    uv run --extra c3d --extra smpl --extra gmr python -m musclemimic.web_viewer.run \
      --c3d-file path/to/file.c3d \
      --c3d-model-path /path/to/smplx/models \
      --c3d-dataset-name MyStudy/Subject01/Trial01
  11. Browse existing retargeted motions

    main

    The musclemimic.web_viewer.run module can also be used to browse existing motions. You must provide either a specific --motion string or a --dataset-group.

    # View a specific motion
    uv run python -m musclemimic.web_viewer.run --motion "KIT/6/WalkInCounterClockwiseCircle06_1_poses"
    
    # View a dataset group
    uv run python -m musclemimic.web_viewer.run --dataset-group AMASS_LOCOMOTION_DATASETS
  12. Use the direct SMPL retargeting backend

    main

    The direct SMPL retargeting backend is used when you provide both --c3d-model-path (for fitting) and --retarget-smpl-model-path (for retargeting). This is an alternative to the default GMR method.

    uv run --extra c3d --extra smpl python -m musclemimic.web_viewer.run \
      --c3d-file path/to/file.c3d \
      --retargeting-method smpl \
      --c3d-model-path /path/to/smplx/models \
      --retarget-smpl-model-path /path/to/retarget/smplh/models