Design

[talmo]

Overview

Goal: Reproduce STAC (PDF) using MJX for GPU acceleration.

This serves the purpose of converting 3D pose tracking into MuJoCo biomechanical model joint angles for subsequent visualization and imitation learning.

Our current implementation of STAC relies on a couple of pieces that we'd like to get rid of:

1. CPU-based physics. This is slow.
2. Uses older MuJoCo (pre 2.0). This is not future proof.
3. Depends on dm_control. This is not future proof and is likely to get deprecated in favor of a Brax-based framework (MuJoCo 3 google-deepmind/mujoco#1101 (comment)).

Since MJX is brand new (as of Oct 2023), this re-implementation of STAC also serves as a trial run for using MJX.

Components

1. Data loaders. We should support both DANNCE outputs (classic and new PyTorch version), as well as SLEAP-Anipose.
2. Rat and mouse models. We should support both the classic rat model from dm_control, as well as our in-development mouse models.
   • Not top priority, but we should be able to support the new fly model from DM/Janelia as well.
   • Maybe also the conventional humanoid one as well?
   • We don't need muscles for STAC technically, so we can just use joint-based actuators.
3. Inverse kinematics and optimization. This is the heavy lifting step, which alternates between solving for offsets between pose landmarks and model joints, as well as model joint angles.

Conceptual issues

1. Do we really not need muscles?
   • Technically we can still use joint angles as a target for imitation, whether or not the actual action space is in muscle activations or joint torques -- but no empirical evidence for this?
2. Will MJX really speed up STAC?
   • We can already trivially do data parallelism over CPUs for STAC.
   • We should be able to treat different (sub)clips as independent "scenes", which MJX parallelizes well over.
   • How do we handle edge effects at the borders of subclips?
   • Is the speed up of processing many clips on a single GPU greater than doing it over many CPUs?
   • Will we have issues with many branching contacts?

[diegoaldarondo]

Clarifications

Compatibility:

• The development branch of STAC is fully compatible with MuJoCo 2.x.

Performance Issues and Possible Solutions:

• The primary issue is not the slowness of CPU-based physics. Instead, it's the serial operation of STAC. Currently, we address this by using multiple processes across thousands of cores.
• A potential solution is to run STAC on multiple clips in parallel on a single node, utilizing GPU acceleration.

Alternating Optimization:

• The part of the process that doesn’t require acceleration is the alternating optimization. This step is a component of the fitting process. It's generally applied to short clips and takes a few minutes at most.
• The offset optimization could be implemented with mjx, but this is not expected to yield significant speed gains.
• The crucial part that needs acceleration is the pose optimization, which is applicable to arbitrarily long clips.

Modifications Needed:

1. Loss Function for Pose Optimization:
   • Requires modifications to make forward kinematics calls.via mjx.
   • Access the offset marker sites via mjx.
2. Optimizer:
   • Currently, it uses scipy.optimize.least_squares.
   • A change to jax.scipy.optimize.minimize or similar is conceivable.
3. Offset Loss Functions:
   • Optionally, reimplement the corresponding functions for the offset loss.
4. There are several dm_control dependencies that are unnecessary and could be simplified while implementing in mjx.

Other

Body Models:

• STAC is configured for models with a body, head, and limbs.
• It should theoretically work for the mentioned models, although tolerances may vary.

Handling Edge Effects at Subclip Borders:

• Edge effects haven't posed significant issues. Solutions tend to be smooth along the boundary.