abdul004/pi05_so101_multitask_checkpoint

SO-101 Pi0.5 Multitask Inference and Compression Journey

This model card captures the path from a single-task SO-101 Pi0.5 policy to a language-conditioned multitask inference setup, then into action-preserving visual compression for remote robot control.

1. From Single Task to Multitask Pi0.5

The project started with a Pi0.5 fine-tune for the SO-101 arm using dual camera observations:

overhead image + wrist image + joint state + language prompt -> Pi0.5 -> action chunk

The first successful policy focused on a single behavior: pick up an object and put it into a container. That established the core remote-inference stack:

• SO-101 follower arm controlled from the Mac mini.
• Overhead and wrist cameras captured through LeRobot/OpenCV.
• Pi0.5 served remotely from RunPod using OpenPI's websocket policy server.
• The Mac client sent image observations and received chunks of continuous joint actions.
• Saved traces included frames, states, actions, action chunks, latency, metrics, GIFs, and start poses.

The multitask step changed the training target from one fixed task string to per-episode language instructions. In OpenPI config terms, the important distinction was:

prompt_from_task = True
default_prompt = None

That allows the model to learn the natural-language task field in the LeRobot dataset instead of conditioning every episode on one hard-coded prompt.

Multitask Variations Tested

The multitask checkpoint was evaluated with several object/container prompts rather than a single fixed "ball in cup" instruction. The point of these runs was to check whether Pi0.5 could bind language to the visible object and target container while using the same SO-101 hardware and camera setup.

Whistle to Green Bin

Successful live RunPod Pi0.5 inference rollout from a fixed start pose:

Additional Multitask Rollouts

The same checkpoint was also tested on other object and target variations.

Lego/Block to Yellow-Blue Target

Tests block-like object geometry and color grounding:

2. Direct Remote Inference

The first goal was simply to run the multitask Pi0.5 policy remotely while controlling the SO-101 arm live:

Mac cameras -> JPEG image payload -> RunPod Pi0.5 server -> action chunk -> SO-101

The direct server is launched with:

uv run scripts/serve_policy.py --port 8000 \
  policy:checkpoint \
  --policy.config=pi05_so101 \
  --policy.dir=/workspace/checkpoint/4999

On the Mac, experiments/policy_trace/record.py records a full closed-loop trace while controlling the robot. For a clean run, we first move the robot to a fixed start pose:

PYTHONPATH=. python experiments/policy_trace/move_to_start_pose.py \
  traces/trace_pi05_whistle_to_bin_runpod_10/start_pose.json \
  --max-relative-target 15 \
  --tolerance 2.0 \
  --sleep-s 0.05 \
  --max-steps 600 \
  --yes

Then run direct inference:

PYTHONPATH=. python experiments/policy_trace/record.py \
  --server "wss://YOUR-POD-8000.proxy.runpod.net" \
  --prompt "Put the whistle in the green bin" \
  --trace-name "direct_whistle_green_bin_runpod_live" \
  --output-dir "traces" \
  --fps 30 \
  --actions-per-inference 10 \
  --max-steps 1200 \
  --max-relative-target 15 \
  --jpeg-quality 80 \
  --frame-quality 92 \
  --start-pose "saved_start_pose"

This answered the first practical question:

Can the current multitask policy do the task from this camera setup and prompt?

Once the answer was yes, the next bottleneck was bandwidth. The policy could run on RunPod, but every inference call still required shipping two camera images from the Mac mini. JPEG made that practical, but it is still a hand-engineered image codec that spends bandwidth on details Pi0.5 may not need and can discard details it may care about. That is where the compression work entered the project.

The next research question became:

Can we send a learned, action-preserving representation instead of camera images?

This was inspired by comma.ai's learned-feature framing for driving: train a representation that preserves control-relevant structure, not necessarily human-perfect pixels. In their planner write-up, comma.ai describes compressing image information into learned features, then using reconstructions to inspect whether planning-relevant details like lane and lead-car locations survived while irrelevant colors/background details could be discarded (comma.ai, "Towards a superhuman driving agent").

For SO-101, the analogous question is:

Can a compact learned representation preserve the object, gripper, container,
and contact geometry that Pi0.5 needs to choose the same action chunk?

That question is what led to the trace decoder, global-latent autoencoder, and finally the spatial-latent live proxy.

3. Raw Action Replay

After a live policy trace succeeds, we can replay the saved actions without calling Pi0.5 again:

saved action chunks -> SO-101

This is intentionally treated as risky because it is open loop. It is useful for separating:

• perception/inference issues,
• latency issues,
• and low-level execution issues.

If a raw action replay from the same start pose behaves similarly to the original trace, the saved actions are physically executable. If replay diverges, then the problem is likely start pose mismatch, object placement mismatch, or open-loop instability.

4. First Compression Experiments: Per-Trace Decoder

The first compression experiments were not a real live codec. They were per-trace reconstruction baselines:

latent_t -> tiny decoder -> generated frame_t

This tested a simple question:

Can generated frames preserve Pi0.5 actions, even if they are imperfect visually?

This is the robot-manipulation version of the comma-style idea. The goal was not to maximize perceptual video quality. The goal was to preserve enough task geometry for the downstream policy to behave the same way:

compression is successful only if Pi0.5 actions stay stable

The action comparison loop used saved robot state from a successful trace and sent either original or generated frames back through Pi0.5:

state_t + generated images_t + prompt -> Pi0.5 -> predicted action chunk

This stage did produce useful robot behavior. The important lesson was more specific: pixel metrics and visual similarity were not enough by themselves. A reconstruction can look plausible, preserve enough information for the task to partly succeed, and still amplify action jerk or change gripper/arm decisions. That moved the project toward action-preservation metrics:

• normalized action delta,
• arm cosine similarity,
• gripper open/close agreement,
• jerk amplification,
• and eventually real robot rollouts.

The decoder-only robot replay is a good example: it showed the idea could work, but it also exposed the high-jerk behavior that made action-level evaluation necessary.

5. Real Encoder-Decoder: Global Latents

The next step was a true autoencoder:

overhead + wrist images -> encoder -> compact latent -> decoder -> reconstructed images

The first version used a single global vector latent. It was compact, around the "1 KB class" depending on dimension and precision, but it discarded too much spatial structure. It could reproduce broad colors and rough layout, but small task-critical geometry suffered:

• gripper shape,
• small object boundaries,
• bin/cup edges,
• and contact details.

This was too brittle for Pi0.5 because the policy depends on spatial relationships, not just scene category.

The visual failure mode was obvious once we looked frame by frame: the global vector kept rough colors and scene mass, but object/contact geometry collapsed into soft blobs.

6. Spatial Latent Autoencoder

The stronger model moved from one global vector to a 2D spatial latent map:

image pair -> C x H' x W' latent map -> image pair

The best working candidate in this phase was:

so101_spatial_lc8_ds8_ch96_320x240

Key idea:

• latent_channels=8
• latent_downsample=8
• base_channels=96
• input reconstruction resolution 320x240

This preserved spatial layout much better than the global vector. Wrist reconstructions became strong enough to preserve the whistle/bangle shape in many frames. Overhead was still weaker: bin edges, gripper/object detail, and small high-frequency cues became soft.

Representative reconstruction frames:

Bandwidth tradeoff:

• uint8 spatial latent payload: about 19 KiB per inference image pair.
• direct JPEG pair at quality 80: much larger and scene-dependent.
• float16 spatial latent: about 2x the uint8 latent payload.

7. Live Latent Inference

The live latent architecture inserted a proxy between the Mac and Pi0.5:

Mac cameras
  -> local autoencoder encoder
  -> quantized latent payload
  -> RunPod latent_decode_policy_server.py
  -> autoencoder decoder
  -> normal Pi0.5 image observation
  -> Pi0.5 action chunk
  -> Mac
  -> SO-101

This avoided modifying the OpenPI Pi0.5 websocket server. The proxy behaved like a man-in-the-middle:

• external client sends latents,
• proxy decodes them to overhead/wrist images,
• proxy forwards normal image observations to Pi0.5,
• proxy returns Pi0.5 actions unchanged.

The first live runs proved the infrastructure worked:

• Mac-side encode after warmup was roughly 10-14 ms.
• RunPod decode was roughly 20-25 ms.
• Pi0.5 policy inference remained around 100 ms inside the proxy.
• End-to-end inference calls were usually around the same order as direct inference.

However, the robot behavior was worse than direct JPEG inference.

8. What Failed in Live Latent Runs

The key symptom was not a network crash. The system moved the arm, but the action stream became less reliable.

Safety clamp warnings were the strongest diagnostic signal. They mean LeRobot refused to send the exact target that Pi0.5 requested because the requested joint target was too far from the current state or outside configured safety limits.

Example:

original goal_pos: -91.3
safe goal_pos: -68.2

In the latent runs:

• shoulder lift and elbow flex were clamped much more often,
• action jerk increased,
• and gripper events were less task-appropriate.

In a cleaner direct JPEG run, clamp events were rare and small. That suggests the latent image path changed the policy's action choices, not just the low-level controller.

Visual inspection matched this:

• wrist decoded images were often acceptable,
• overhead decoded images were only "so-so",
• bin edges and small object/gripper structure were softened,
• and the policy likely lost important spatial cues.

The comparison below shows local camera frames against the server-decoded latent frames from a live latent rollout: