# modbot — simulation and training for the modular legged robot

Everything the paper's sim side needs, runnable today on CPU, fast on one GPU.

```
generator.py   module graph -> MJCF; 10-body physical test set; random body sampler
env.py         MJX environment; one batched model, a different BODY per env
train.py       PPO (brax) over a population of bodies; the scaling sweep
evaluate.py    run a checkpoint on all 10 test bodies -> sim side of the scaling curve
compat.py      jax 0.11 / brax 0.14 shim
```

## Install
```
pip install mujoco mujoco-mjx brax
# GPU: pip install -U "jax[cuda12]"
```

## The trick that makes multi-embodiment training cheap
MJX batches environments only if they share one model. So every env uses the
same "master" hexapod model; a body with fewer legs or different link lengths is
the master with different *parameters*: masked legs become near-massless, lifted
inside the plate, with zero actuator gain. The randomization hook assigns one
sampled body per env. The policy sees a 6x5 morphology block
[present, cos, sin, thigh, shin] so it knows what it is driving. This is the
same observation the real robot will produce from its self-description.

## Run
```
python3 generator.py                 # 10 test bodies stand under PD control
python3 train.py --smoke             # CPU pipeline check, learns nothing
python3 train.py --mode scaling --n_bodies 10   --steps 100_000_000 --tag n10     # GPU
python3 train.py --mode scaling --n_bodies 1000 --steps 100_000_000 --tag n1000
python3 evaluate.py runs/n1000.pkl   # sim scaling-curve data for the 10 test bodies
```
The sweep is `--n_bodies 1 / 10 / 100 / 1000` with the same seed; plot
`success` from each `*_eval.json` against n_bodies. That is Figure 1 (sim); the
robot gives the matching real line.

## What is placeholder and must be replaced
- `SERVO_KP_DEFAULT`, damping, armature, frictionloss in generator.py: replace with the
  identified STS3215 12 V model (pendulum rig, week 9).
- Foot friction and contact solref: measure the cast silicone feet (week 10).
- Observation noise and latency: add once the real bus loop rate is measured.
- Policy architecture is a padded MLP. A URMA-style per-joint token policy is the
  upgrade once the baseline trains; the env's obs layout already separates per-leg blocks.

## Known caveat
brax prints that its env wrappers are in maintenance mode and points to
mujoco_playground. The env here is plain MJX; porting it to mujoco_playground or
mjlab later is a wrapper change, not a rewrite. Newton (NVIDIA's Warp engine) also
loads MJCF, so the generator is reusable there too.