feat: add Drake backend via DrakeUni

[Langford-H]

Summary

This is a large Drake backend feature update. Please review cautiously: it adds a new simulator backend path, an external DrakeUni batch runtime dependency, backend tests, replay support, and initial Drake task configs.

Important: to use or validate sim_backend=drake, reviewers must install Drake C++ first. DrakeUni is built against that Drake install and is required by UniLab's Drake backend. Drake is absent only from non-Drake CI/import paths; the Drake backend itself is not usable without Drake + DrakeUni.

• Add a Drake backend path through the external DrakeUni batch runtime.
• Add initial Drake task configs:
  • PPO: go1_joystick_flat/drake, go2_joystick_flat/drake, go2_footstand/drake, go2w_joystick_flat/drake, stewart_balance/drake
  • SAC/off-policy: go2_joystick_flat/drake, go2_footstand/drake, go2w_joystick_flat/drake, stewart_balance/drake
• Add Drake replay support where Drake performs physics stepping and MuJoCo renders recorded videos.
• Keep non-Drake workflows isolated from Drake dependencies: normal upstream CI does not install Drake/PyDrake/DrakeUni, and Drake runtime tests are gated when DrakeUni is unavailable.
• Align the boundary with MuJoCoUni's state-first design: DrakeUni owns compact batch stepping and raw MJCF-declared sensor output; UniLab owns task-facing sensor/body views and SimBackend adaptation.

Scope:

• Go1/Go2 joystick-flat are the main reproducible train/replay smoke tasks.
• Additional committed configs broaden the initial Drake surface, but should be treated as early validation targets rather than full Drake/MuJoCo parity claims.
• No generated checkpoints, replay videos, TensorBoard logs, or training output directories are intended to be committed.

Linked Work

• Issue:
• Milestone: Drake backend foundation / DrakeUni integration
• DrakeUni runtime repo: https://github.com/Langford-H/drakeuni
• DrakeUni benchmark repo: https://github.com/Langford-H/drakeuni_bench

Review Guidance

Suggested review order:

1. Verify existing MuJoCo/Motrix paths are not behaviorally changed.
2. Verify the dependency boundary:
   • UniLab imports without Drake when the Drake backend is not used.
   • Upstream non-Drake CI does not require Drake/PyDrake/DrakeUni.
   • Drake runtime tests skip cleanly when DrakeUni is absent.
3. Review the Drake backend contract:
   • DrakeUni remains task-blind and state-first.
   • UniLab owns named sensor/body views.
   • Drake replay steps with Drake physics, not MuJoCo physics.
4. Review task configs separately from backend machinery.
5. Treat policy quality as early validation only; this PR is about backend integration.

Validation

• make check
• uv run pytest -m "not slow"
• Additional Drake-specific validation listed below.

Final validation result:

make check
passed

uv run pytest -m "not slow"
1307 passed, 27 skipped, 258 deselected, 68 warnings in 35.02s

Reviewer steps for Drake runtime validation:

# 1. Install Drake C++ first.
# Set DRAKE_HOME to a Drake install prefix containing:
#   $DRAKE_HOME/include/drake
#   $DRAKE_HOME/lib/libdrake.so
export DRAKE_HOME=<DRAKE_HOME>

# 2. Clone and install DrakeUni.
git clone https://github.com/Langford-H/drakeuni.git /path/to/drakeuni
cd /path/to/drakeuni
uv pip install -e .

# 3. Build the DrakeUni C++ extension against Drake.
uv run python scripts/build_drake_batch.py --drake-home "$DRAKE_HOME"

# 4. Install DrakeUni into the UniLab environment.
cd /path/to/UniLab
uv pip install -e /path/to/drakeuni

# 5. Run focused Drake tests from UniLab.
PYTHONPATH=/path/to/drakeuni/src \
  uv run pytest tests/base/backend/test_drake_batch_pool.py

Additional commands run during development:

# Config/script tests covering Drake task composition.
PYTHONPATH=/Users/huanghaochen/solver/drakeuni/src \
  uv run pytest tests/scripts/test_train_scripts.py -k "drake or go2"

# Go1 Drake minimal train smoke.
HYDRA_FULL_ERROR=1 PYTHONPATH=/Users/huanghaochen/solver/drakeuni/src \
  uv run python scripts/train_rsl_rl.py \
  task=go1_joystick_flat/drake \
  training.no_play=true \
  algo.num_envs=4 \
  algo.num_steps_per_env=4 \
  algo.max_iterations=1 \
  env.drake_nthread=2

# Go1 Drake replay smoke.
HYDRA_FULL_ERROR=1 PYTHONPATH=/Users/huanghaochen/solver/drakeuni/src \
  uv run python scripts/train_rsl_rl.py \
  task=go1_joystick_flat/drake \
  training.play_only=true \
  training.play_render_mode=record \
  training.play_steps=120 \
  training.play_env_num=1 \
  algo.load_run=<go1_drake_run> \
  env.drake_nthread=1

# Go2 Drake minimal train smoke.
HYDRA_FULL_ERROR=1 PYTHONPATH=/Users/huanghaochen/solver/drakeuni/src \
  uv run python scripts/train_rsl_rl.py \
  task=go2_joystick_flat/drake \
  training.no_play=true \
  algo.num_envs=64 \
  algo.num_steps_per_env=8 \
  algo.max_iterations=5 \
  env.drake_nthread=20

# Go2 Drake replay smoke.
HYDRA_FULL_ERROR=1 PYTHONPATH=/Users/huanghaochen/solver/drakeuni/src \
  uv run python scripts/train_rsl_rl.py \
  task=go2_joystick_flat/drake \
  training.play_only=true \
  training.play_render_mode=record \
  training.play_steps=120 \
  training.play_env_num=1 \
  algo.load_run=<go2_drake_run> \
  env.drake_nthread=1

Several additional local train/replay rounds were run while developing the Drake backend and DrakeUni parser/runtime boundary. These helped uncover generic MJCF parser/runtime issues, including contact handling, geom/materialization edge cases, inherited actuator defaults such as kp/kv, and generated asset visibility. Generated checkpoints/videos/logs are not committed.

Impact

• Backend impact: adds drake; existing mujoco / motrix paths intended unchanged.
• Platform impact: core UniLab CI remains Linux-compatible without Drake; Drake path was manually validated with a local Drake C++ install.
• Training effect expected: yes, for new Drake task targets only.

Artifacts

• W&B:
• benchmark result: https://github.com/Langford-H/drakeuni_bench
• video / screenshot: local replay MP4s were generated during validation, but are not committed
• ONNX / checkpoint: local Drake smoke checkpoints were generated during validation, but are not committed

Checklist

• Added or updated tests where needed.
• Updated docs if behavior or workflow changed.
• Linked the driving issue.
• Noted follow-up work explicitly.

Follow-up work:

• Add a Drake/DrakeUni CI job once upstream maintainers agree on the Drake install path.
• Expand Drake task coverage beyond the currently validated tasks.
• Continue DrakeUni parity work for additional MJCF actuator/sensor semantics.
• Continue performance profiling and optimization after the generic backend boundary is reviewed.

[TATP-233]

NB！容我们仔细看看