A1X wrist-camera grasp policy (sim)
The last centimetres of a grasp for the Galaxea A1X arm with its G1 parallel gripper, trained in MuJoCo. A planner (or a VLA) brings the gripper to a hover a few centimetres above a proposed top-down pinch grasp, with the error real perception leaves behind. From there this policy owns the motion until the object is lifted: small tool-centre-point steps, a yaw step, and the gripper target, 20 times a second. A failed close is not the end: the policy opens, re-aims and tries again.
Two policies are in this repository.
| file | what it sees | how it was trained | success, 100 held-out episodes |
|---|---|---|---|
student.pt |
wrist camera image + proprioception | imitation of the teacher (DAgger, 3 rounds) | 79 % hard mix, 94 % plain objects |
teacher.pt |
the true object pose (simulator state) | PPO | 94 % hard mix, 96 % plain objects |
"Plain objects" are random boxes, cylinders, spheres and capsules. The "hard mix" adds tools with heads, pens, L and T shapes, pucks and flat bars, shoves the object 1 to 2 cm during the descent in 30 % of episodes, and forces the first close early in another 30 %, so the policy must recover.
Status: simulation only. Nothing here has moved the real arm yet. The wrist camera pose
and lens the student was trained with are the ones measured on the real arm on
2026-09-26 (wrist_camera_spec.json), so the next step is the student on real wrist frames.
Student inputs and outputs
Inputs, every tick:
- the wrist camera frame, 160 x 120 RGB, from the camera pose in
wrist_camera_spec.json(frame is 4:3; a 16:9 real frame is centre-cropped first); - a fourth image channel: a Gaussian blob at the pixel where the believed grasp point
projects through the calibrated camera (
env.believed_uv()in the source); - a 19-vector: the 6 joint angles normalised to [-1, 1] by their ranges, the fingertip gap (m x 10), the believed grasp point relative to the tool centre point (m x 10, 3 values), sin and cos of twice the believed yaw error, the current gripper target in [0 closed, 1 open], the previous action (5), and the number of close attempts so far divided by 3.
Output: 5 values in [-1, 1]: TCP step in x, y, z (times 8 mm), yaw step (times 5 degrees), gripper target (mapped to closed at -1, open at +1). The base frame has the table at +z up; the tool centre point frame has x as the approach and y as the jaw closing axis.
Architecture: ImageNet-pretrained ResNet18 trunk with a 4-channel first convolution, a small MLP on the proprioceptive vector, a 3-layer head with tanh. 11 M parameters.
Loading
import torch, numpy as np
from student import Student # student.py in this repository
st = Student.load("student.pt") # device="cuda" if you have one
action = st.act(image[None], obs[None], uv[None])[0] # image uint8 (120,160,3), obs float32 (19,), uv (2,)
The teacher is a plain MLP actor-critic (Policy in sim/grasp/ppo.py of the source
repository) over a 27-value privileged state; it is provided for further DAgger rounds.
Training recipe
Environment and code: sim/grasp/ in the source repository (MuJoCo 3.13, gymnasium,
torch 2.11). Teacher: PPO, 8 to 24 parallel environments, ~1 M steps, warm-started from a
primitives-only run. Student: 4000 episodes per DAgger round, teacher labels on every tick,
the student driving 50 / 70 % of ticks in rounds 1 / 2, 10 epochs per round on an H100,
Huber loss, crop and photometric augmentation, per-scene wrist camera jitter of 3 mm / 0.7 deg.
Known limits, measured: the student's remaining failures are episodes where the object leaves the wrist camera's view after a shove or a failed close (the teacher, which knows the true pose, still succeeds). Thin pens (10 to 18 mm) fail for both policies because the planner's grasp proposal keeps the fingertips 6 mm above the table.
Source
Code, environment, calibration data and the box payloads: machinekind/galaxeo-manipulators,
branch grasp-rl, directory sim/grasp.