DrakeUni

简体中文

Experimental Drake batch simulation runtime for UniLab's Drake backend adapter.

Current scope:

• Compiled C++/pybind batch pool driven by parsed MJCF model contracts.
• nbatch and nthread worker control.
• Python API modeled after MuJoCoUni's BatchEnvPool direction.
• Intended to be consumed by UniLab's Drake backend adapter.

This is not a full Drake backend yet. The current runtime expects a model file, environment count, simulation step size, and worker count. Task semantics such as base-body meaning, named sensor views, rewards, and observations belong to the consuming UniLab backend/task layer.

The input model is the original UniLab MJCF scene, for example scene_flat.xml. DrakeUni materializes a temporary Drake-compatible copy at runtime by expanding MJCF defaults and converting unsupported mesh formats when needed. Callers should not commit or pass pre-generated files such as *_drake.xml or converted OBJ asset directories.

Directory Roles

src/drake_uni/runtime/ is the Python-facing runtime interface. It owns the public config/data contracts, MJCF contract parsing, Drake-compatible MJCF materialization, runtime construction, and the reset/step/sensor API used by UniLab.

src/drake_uni/compiled/ is the native extension layer. It contains the C++ Drake batch pool source and the built pybind extension that performs batched physics stepping.

The intended call path is:

UniLab DrakeBackend
  -> drake_uni.runtime
      -> drake_uni.compiled.DrakeEnvPool
          -> Drake C++ simulation

Build Batch Extension

uv run python scripts/build_drake_batch.py --drake-home /Users/huanghaochen/solver/drake/install

The extension is written to src/drake_uni/compiled/_drake_env_pool*.so. The Drake prefix must provide include/drake, include/pybind11, and the Drake shared library. The generated extension is local build output and is not committed.

Development

The repository uses uv for reproducible local development:

make sync
make install
make check
make build DRAKE_HOME=/path/to/drake/install

make check runs Ruff and pytest. Parser tests that depend on a sibling UniLab checkout skip themselves when its assets are unavailable. See AGENTS.md for the native build contract, scope boundaries, and release workflow.

Install Editable

From a consuming project:

uv pip install -e /Users/huanghaochen/solver/drake_uni

Runtime API

from drake_uni.runtime import DrakeBatchConfig, create_runtime

runtime = create_runtime(
    DrakeBatchConfig(
        model_file="/path/to/scene_flat.xml",
        num_envs=32,
        sim_dt=0.002,
        nthread=8,
    )
)

The preferred integration point is drake_uni.runtime. DrakeEnvPool and the compiled extension are lower-level implementation details.