tiantianx/aloha2-smolvla-banana-apace-d20-jepa-scratch
A-PACE ALOHA banana: scratch JEPA, d=20
Interim JEPA-only checkpoint: 12,600 optimizer steps (epoch 10/20). This is a new, independently initialized JEPA with balanced visual, physical-joint, and gripper prediction losses. The SmolVLA policy remains frozen. The original training budget is 25,200 steps. This release is a fixed snapshot, not a statement that training has finished.
No adapter is included. Omit `--aca-checkpoint`; do not pair this JEPA with an adapter from an earlier repository. Matching adapter training and evaluation are pending.
Pinned dependencies
- Base policy:
yhong96/aloha2_smolvla_banana_v2, revision558170cf34a8a5c83e5ce68944846186a82a4b80. - Dataset:
yhong96/aloha2_banana_insertion, revision213b3da1b40e35ab67278eafe2415461982ce65e. - Supported nonzero delay: 20 control steps, nominally 400 ms at 50 Hz.
- Native action horizon: 50. Action coordinates:
delta_joint_v1, anchored to the predicted handover joint state. - Committed inputs: 20 absolute 14D ALOHA command vectors, 12 arm joints plus 2 grippers, in the recorded channel order.
Use with the existing ALOHA A-PACE rollout package
from huggingface_hub import snapshot_download
snapshot_download(
repo_id="tiantianx/aloha2-smolvla-banana-apace-d20-jepa-scratch",
local_dir="models/apace_d20_scratch",
)Set the existing policy server option:
--jepa-checkpoint models/apace_d20_scratch/jepa_training/best.ptLeave --aca-checkpoint unset (remove any old adapter argument in your launcher). Use the pinned base policy and its native processors, as verified by apace_model_binding.json. This repository contains predictor weights, not the base policy or a standalone robot server.
Supply the actual committed absolute commands. At the scheduled handover, the future-conditioned action chunk starts at index 0; do not skip 20 actions again. JEPA predicts the native normalized joint state. The existing rollout runtime restores absolute arm targets by adding the generated joint offsets to the predicted handover state; gripper outputs remain absolute. Do not apply this conversion twice. These weights do not change the rollout scheduler or controller.
Held-out prediction validation
Same 8 held-out episodes / 2,240 windows as the old 25,200-step predictor. These are future-state prediction errors, not executed-trajectory jerk.
Joint prediction improves substantially, while visual feature MSE is still higher. Full-sampler action evaluation and real-robot closed-loop performance have not been established for this snapshot. These offline results do not establish robot success, physical smoothness, or inference latency.
Training
Full JEPA random initialization; cached frozen-policy condition features and normalization reused. Batch 16 (microbatch 8, accumulation 2), AdamW peak LR 2e-4, OneCycle schedule, 72 training episodes and 8 validation episodes. Loss: visual MSE / train-only scale + physical arm-joint MSE / train-only scale
- gripper normalized MSE / train-only scale + 0.1 cosine distance. Outputs retain the original native normalized-state interface.
The run was interrupted after step 12,525 and restored from its step-12,000 checkpoint, including optimizer and scheduler. That checkpoint did not contain RNG state, so the recovery is not bitwise identical to an uninterrupted run. This does not change the fact that the run originally started from random initialization.
jepa_training/best.pt is selected by held-out physical arm MSE at snapshot time. See CONFIG.json, VALIDATION.json, RELEASE.json, and SHA256SUMS.json for configuration, raw metrics, step identity, and file integrity.
