mjfit

Exact convex collision geometry for MuJoCo pegs, holes, threads and gears

License

MIT

Overview

mjfit is a Python library that generates collision geometry for assembly tasks in MuJoCo: a peg and its hole from any outline, a bolt and its nut, or a pair of meshing gears, each built from convex pieces whose union is exactly the part. It comes from Pulkit Agrawal's Improbable AI Lab at MIT, where it was written to train the policies in the paper Blind Precise Assembly: Proprioceptive Sim-to-Real Peg-in-Hole for Diverse Parts, and it is released under the MIT licence on GitHub.

The problem it solves is specific. MuJoCo collides only convex shapes, and a hole is concave. The usual fixes are approximate convex decomposition, with CoACD or V-HACD, or MuJoCo's signed-distance (SDF) collision on a mesh. Decomposition fills in the very clearance an assembly task depends on, and SDF collision is slower and misses contacts. mjfit instead cuts the plate around a hole into convex wedges, makes the peg the hole offset inward by the clearance, and sweeps thread and gear teeth from convex profiles, so every peg edge sits exactly one clearance from its hole edge and MuJoCo collides every piece natively. Curves become polygons that circumscribe them, so a round hole is never smaller than drawn.

Each generator returns an MJCF string, from Python or from the mjfit command line, and outlines and dimensions are parameters, so one training run can hold thousands of different assemblies. The project page runs the paper's policies and mjfit's comparisons live in the browser.

mjfit has three generators. mjfit.peg takes the outline of either the hole or the peg, as a list of vertices or a shape from mjfit.shapes (circle, rectangle, stadium, trapezoid, regular polygon, star, any polygon, Bezier outlines with rounded corners, or a CSV file), and builds the other as its exact offset by the clearance. Options add material standing inside the hole, such as a USB tongue or connector pins (the peg becomes a hollow shell around them), a wider grip section, a convex outer plate with a minimum rim, a closed hole, or a fixed peg mass. mjfit.thread builds a bolt and its nut from ISO metric, trapezoidal, ACME, square, buttress, round-wire or custom tooth profiles, single- or multi-start and left- or right-handed. mjfit.gear builds a spur gear or a meshing pair with involute teeth on hinges, with an optional motor; each tooth is one convex prism.

Configurations are YAML files holding a generator's arguments, and the package ships ready ones: circle, star, USB-A, USB-C, HDMI, RJ45, D-sub (DB9) and XLR3 pegs, M10 and M6 threads and a FurnitureBench leg, and several gear pairs. For training, mjfit.sample_pegs draws reproducible random assemblies from two families (random convex outlines with rounded corners, or primitive shapes with up to three islands in the hole), and mjfit.grid lays many assemblies out in one scene with names that never clash.

The output is ready to use: every geom belongs to a default class with stiff contacts, files can run on their own or be included in a larger scene, and meshes can be written as OBJ files when a USD converter needs them. The repository's benchmark compares mjfit with CoACD, V-HACD and MuJoCo's mesh SDF on a 16 mm round or square peg with clearances down to 20 µm, on CPU and on GPU with MuJoCo Warp, and ships the commands that reproduce every number.

Why it matters

Contact-rich assembly is one of the hardest things to learn in simulation, and the hardest part is often not the learning but the geometry: if the simulated hole is smaller, rougher or slower than the real one, a policy trained on it learns the wrong lesson. mjfit makes exact, fast, procedurally varied pegs, holes, threads and gears a pip install away for anyone using MuJoCo, including MuJoCo Warp for training at scale on a GPU. The paper it was built for shows what that enables: blind insertion policies trained entirely in simulation that transfer to a real arm and insert pegs of shapes they never saw.

Problem it solves

MuJoCo only collides convex shapes, so a plate with a hole has to be approximated. Convex decomposition closes up the clearance that a peg-in-hole, threading or gear-meshing task depends on, and signed-distance collision is slow and misses contacts. mjfit builds each part from convex pieces whose union is the exact part, with the clearance kept, and generates thousands of varied assemblies from parameters.

Key takeaways

• Builds MuJoCo pegs and holes, bolts and nuts, and meshing gears from convex pieces whose union is the exact part
• Every peg edge sits exactly one clearance from its hole edge; circles become circumscribed polygons, so holes are never smaller than drawn
• On CPU a peg and its hole step an order of magnitude faster than MuJoCo's mesh SDF at equal or better accuracy, per the repository's benchmark
• Bundled configs for real connectors and fasteners: USB-A, USB-C, HDMI, RJ45, D-sub, XLR, M10 and M6 threads, gear pairs
• Procedural sampling draws thousands of reproducible assemblies for a single training run
• Built for the MIT paper Blind Precise Assembly, whose insertion policies transfer from simulation to a real arm

Who it is for

• Researchers training assembly, insertion, threading or gear-meshing policies in MuJoCo or MuJoCo Warp
• Robot-learning teams who need many varied parts with tight, exact clearances
• Simulation engineers who have hit the limits of convex decomposition or SDF collision on concave parts
• Students learning contact-rich manipulation who want ready-made connector and fastener models

Limitations

• Early release: version 0.1.0, maintained by one author, and installed from GitHub rather than PyPI
• MuJoCo only; other simulators get the meshes as OBJ files, not the contact model
• Generates parametric parts (pegs and holes, threads, spur gears), not arbitrary CAD models
• Stiff contacts need a time step of 1 ms or less, and hinged gears need MuJoCo's PGS solver rather than the default Newton solver
• Thread wedges lose a little material to the helix's discretisation, so threads are exact only to that tolerance