You can deploy this model and start earning money today!

Ripito/
pickplace-mjx-day11-weld

Name: pickplace-mjx-day11-weld
Availability: InStock

jax

reinforcement-learning

Views

No views yet

pickplace_mjx — SO-101 single-arm pick-and-lift (day11_weld_peak)

A Brax PPO policy trained in MuJoCo MJX for the LeRobot SO-101 5-DOF scissor-gripper arm. Task: pick a 4×4×6 cm cube off a tabletop and lift it to a target point in the air above the workspace.

This is the day11 checkpoint of the SOP-VLA project's sim2real pipeline — the FIRST policy in the pipeline that loads cleanly from disk and reproduces its training metrics in external evaluation (prior checkpoints had a normalizer + magic-grasp eval-gap).

Highlights

Real grip physics in sim: uses MuJoCo <weld> equality with soft solimp to model grip-on-cube without the qpos-teleport "magic grasp" hack used in earlier checkpoints. Closer to real-robot behavior.
Eval-gap closed: training reports episode_success=118, external Brax Evaluator reports 97-127 on the saved pickle (matches).
30% binary episode success rate in 100-episode external eval — cube enters the 4cm goal zone in 30% of trials. Median min-distance to goal: 5.8 cm.

Files

File	Purpose
`policy.pkl`	Pickled `(RunningStatisticsState, policy_params, value_params)` tuple from Brax PPO `train_fn` return. Direct load via `pickle.load`.
`network_meta.json`	Network factory kwargs needed to reconstruct the inference function (obs_size, action_size, hidden layer sizes, activation, normalize_observations).
`eval_results.json`	100-episode external eval: episode_success_rate, mean_min_dist, time_to_first_success, per-episode breakdown.

Usage

python

import pickle, json, functools, jax
from brax.training.agents.ppo import networks as ppo_networks
from flax.linen import silu

# Load
with open("policy.pkl", "rb") as f:
    params = pickle.load(f)
meta = json.load(open("network_meta.json"))

# Reconstruct
nets = ppo_networks.make_ppo_networks(
    observation_size=meta["obs_size"],
    action_size=meta["action_size"],
    policy_hidden_layer_sizes=tuple(meta["policy_hidden_layer_sizes"]),
    value_hidden_layer_sizes=tuple(meta["value_hidden_layer_sizes"]),
    activation=silu,
)
inference_fn = ppo_networks.make_inference_fn(nets)(params, deterministic=True)
inference_fn = jax.jit(inference_fn)

# Run on a 32-dim obs vector
action, _ = inference_fn(obs, jax.random.PRNGKey(0))
# action is 6-dim: 5 arm-joint deltas + 1 gripper delta, scaled by 0.03 rad

Training config

Reproduced via python pickplace_mjx/train.py --no-normalize-obs --total-timesteps 60000000 in the VLA-Models repo at commit 4f8df3e or later.

Hyperparam	Value
Algorithm	PPO (Brax)
Architecture	MLP 2×256, silu, lecun_uniform init
Action distribution	tanh-normal, scalar log_std, init=1.0
`num_envs`	1024
`unroll_length`	32
`num_minibatches`	4
`num_updates_per_batch`	5
`learning_rate`	1e-4
`entropy_cost`	0.01
`clip_epsilon`	0.2
`discounting` (γ)	0.99
`normalize_observations`	False (workaround for an Orbax round-trip bug in Brax — see below)
`total_timesteps`	60M env-steps target (overshot via brax counter)
Sim backend	mujoco-warp 1.12 (MJX warp impl)

Environment

SO101PickPlace from pickplace_mjx/env.py:

Robot: 5-DOF SO-101 arm + 1-DOF scissor gripper, MJCF from mujoco_menagerie/robotstudio_so101
Object: 4×4×6 cm rigid cube, 50 g, friction 2 0.05 0.005, condim=6, priority=2
Goal: floating mocap target, position randomized in x ∈ [0.08, 0.18], y ∈ [0.08, 0.22], z ∈ [0.10, 0.14] per episode
Cube spawn: same xy region, z = 0.03 (resting on table)
Episode length: 200 steps at 50 Hz (4 seconds)
Action: 6-dim, ±1, scaled by 0.03 rad/step joint delta
Observation: 32-dim flat vector — arm qpos/qvel (5+5), gripper qpos/qvel (1+1), EE pos (3), cube pos (3), cube quat (4 — placeholder zeros), goal pos (3), cube-EE relative (3), goal-cube relative (3), prev_grasped flag (1)
Reward: dense shaped — reach + pregrasp + held + lift + transport
- near_miss + sparse success (20.0 at <4cm, 0 else) + action_pen, clipped to [-10, 10]

Grip mechanism (the weld)

MuJoCo <weld> equality between moving_jaw_so101_v1 and box, active="false" initially. The env activates data.eq_active[weld_id] when the grasp predicate (gripper_qpos > 0.3 AND ee_to_cube_dist < 0.07) fires. Soft solimp="0.7 0.9 0.001" lets the cube settle into a fixed relpose (0, -0.05, 0.02, identity_quat) smoothly without catapult.

This is closer to real-robot grip behavior than the qpos-teleport "magic-grasp" used in day7-day10 checkpoints, which had a 3-∞× train-vs- eval discrepancy.

Known limitations

Normalizer disabled: Brax's RunningStatisticsState doesn't round- trip cleanly through Orbax checkpoints. Enabling normalize-obs gave the policy training peak success=99-117 but external eval = 0. Using normalize_observations=False slows learning ~10% but produces a serializable policy.
Sim-to-real gap on grip force: real scissor-gripper contact friction is higher than what MJX models. Magic/weld in sim is a stand-in for real grip happening. Real-robot deployment (Phase 2 of the project) uses teleop demos for the grip portion, not a sim policy.
Mean min-distance ≈ 5.8 cm: most failed episodes get ~6 cm from goal but don't enter the 4 cm zone. Precision could be improved with a curriculum on the success threshold (warm-start fine-tune attempt on day12 destabilized the policy — try smaller deltas next time).

License

Apache-2.0, same as Brax + MuJoCo Menagerie + LeRobot. No external training data; weights trained from random init in sim.

Citation

@misc{kitano_2026_pickplace_mjx_day11,
  title  = {pickplace_mjx — SO-101 single-arm pick-and-lift PPO baseline},
  author = {Rihito Kitano},
  year   = {2026},
  url    = {https://huggingface.co/Ripito/pickplace-mjx-day11-weld},
  note   = {Phase 1 sim baseline of the SOP-VLA project, MJX + Brax PPO}
}