AssembleNIST: Assembly Policy Learning in MuJoCo

By Tao Sun | Stanford University

A single reinforcement-learning (PPO) policy for all 20 assemblies of the NIST Assembly Task Board 1 by a Franka arm, learned from binary rewards with a Franka arm in MuJoCo.

This repository is an independent implementation of the approach introduced by AlphaNIST and Success Guided Sampling (SGS). It is built in MuJoCo environment.

Since the AlphaNIST hasn’t yet released its source code (as of Oct 21, 2026), I made own design choices for contact modeling, the action space and the start-state distributions. I am not affiliated with or endorsed by the AlphaNIST authors.

With 8,192 parallel environments, the policy reaches 98.4% evaluation success after 10K PPO iterations, outperforming the AlphaNIST’s performance with less than half of the 25K iterations reported.

Success rate (%) by training iterations.

The full learning curve can be found in results/reference_e8192.csv.

20 parts at random initial poses

After assembly

Success rate (%) by part type.

Success rate (%) by starting state.

More visuals. MP4 files can be found in docs/videos.

Clone the repository, install the dependencies and build the contact plugin (requires g++ on x86-64):

git clone https://github.com/suniique/assembly-nist.git && cd assembly-nist
pip install -e .
bash plugin/build.sh

Generate a demo with the pretrained policy:

MUJOCO_GL=egl python -m assembly_policy.viz.demo --ckpt checkpoints/e8192_it12500.pt --out demo.mp4

Use MUJOCO_GL=osmesa if EGL is unavailable. The demo generates settled start poses if reset banks are absent; --seed changes the layout and assembly order, and --pin 2:N16+3:connector fixes positions in the order. Pretrained policies after 2.5K, 5K, 7.5K, 10K and 12.5K iterations are in checkpoints/.

First, we need to build the reset banks before training:

python scripts/build_banks.py --workers 16

Evaluate the pretrained policy. It runs 3200 episodes for about one hour on 16 CPU cores:

python -m assembly_policy.evaluate --policy checkpoints/e8192_it12500.pt --out eval.json

For training the policy, you can run with one machine on a smaller environment size (e.g., 512):

python -m assembly_policy.train --run-dir runs/small --envs 512 --nthread 32

To achieve best performance, the default training configuration uses 8,192 environments and 10K PPO iterations. Below is a script to train it on a Slurm cluster.

sbatch scripts/slurm/build_banks.sbatch
sbatch --export=ALL,RUN_DIR=runs/e8192,NCONN=960 scripts/slurm/learner.sbatch
for i in $(seq 30); do
  sbatch --export=ALL,RUN_DIR=runs/e8192 scripts/slurm/workers.sbatch
done

My first long runs on 8k environments stalled out at a 46% success rate after 23k iterations. Getting the policy to learn efficiently took several non-obvious changes. Below are the lessons that I feel matter most.

1. Keeping RL exploring helps in this task. With entropy 0.001, an initial std of 0.5 and a std floor of 0.1, the SR crept up about 1 point per 1000 iterations. A learned std starting at 1.0 with no floor and entropy 0.005 reached 49% after 2000 iterations and 98% after 5500.
2. Make the wrist fast enough to finish in 15s. Under the libfranka gains, a relative target $$q + s \cdot a$$ never fully catches up. With a uniform 0.02 rad scale, the wrist took 15–27s just to rotate 90°. Tuning per-joint scales to give every joint ~0.5 rad/s per unit action finally made reorientation possible.
3. Fix the physics before tuning RL. Nut-threading sat at 0% SR because its sharp thread tips were getting stuck in MuJoCo's SDF contact. Adding 0.2 mm of radial clearance let the nuts spin freely without skipping threads. I also had to tune the timestep: gear teeth and threads needed 0.5 ms steps with stiff contacts.

Final design

• Policy: a SimBa-style actor and critic trained with PPO, using part, fixture, and gripper point clouds plus proprioception history.
• Control: relative joint-position targets and a binary gripper at 25 Hz with libfranka impedance gains, per-joint action scales, and an action-rate penalty.
• Reward: binary assembly success at 1 mm keypoint tolerance, with small control penalties. Episodes end on success, a joint-speed violation, or the 15-second limit.
• Start states: loose, pickup, transport, partial, and near-mating poses, sampled with SGS.
• Contacts: primitive geometry for rods and pegs, convex gear teeth, helical-thread SDFs for nuts and bolts, and mesh SDFs for connectors. Gears and nuts use 0.5 ms simulation steps and other parts use 2.5 ms.

This project builds on

• AlphaNIST: Emergent Assembly Skills from Binary Task Rewards.
• A Balanced Data Diet: Addressing the Exploration Bottleneck in Mega-Scale RL for Robot Control (Zhang et al., 2026).

Sincerely thanks to Octi Zhang for grateful guidance with training details (including reset bank configurations, PPO settings, and state sampling) and Jing Zhang for insightful discussions! The NIST board assets come from UW Lab and the robot model from MuJoCo Menagerie.

If you use AssembleNIST, please cite this repository together with AlphaNIST and SGS (all entries are also in CITATION.bib):

@misc{sun2026assemblenist,
  title        = {AssembleNIST: One Policy for All NIST Board Assemblies in MuJoCo},
  author       = {Sun, Tao},
  year         = {2026},
  howpublished = {\url{https://github.com/suniique/assembly-nist}}
}

@article{zhang2026alphanist,
  title   = {AlphaNIST: Emergent Assembly Skills from Binary Task Rewards},
  author  = {Zhang, Octi and Akinola, Iretiayo and Scalise, Rosario and Guaman Castro, Mateo and Yin, Patrick and
             Widmer, Tobias and Macklin, Miles and Amevor, Chris and Turpin, Dylan and Van Gelder, Dirk and
             Carius, Jan and Handa, Ankur and State, Gavriel and Guo, Yunrong and Narang, Yashraj},
  journal = {arXiv preprint arXiv:2610.13354},
  year    = {2026}
}

@article{zhang2026sgs,
  title   = {A Balanced Data Diet: Addressing the Exploration Bottleneck in Mega-Scale RL for Robot Control},
  author  = {Zhang, Octi and Guaman Castro, Mateo and Yin, Patrick and Dagnino, Ignacio and Gupta, Abhishek and
             Scalise, Rosario and Boots, Byron},
  journal = {arXiv preprint arXiv:2610.12465},
  year    = {2026}
}

The code is released under the MIT License. See THIRD_PARTY.md for third-party assets, codes, and their licenses.

For questions and feedback, please reach out to taosun@stanford.edu.