Add per-frame proprio context and JPEG-compress temporal image stacks#465
Conversation
Video-conditioned policies train on real proprio on every clean frame. `TemporalFrameStack` stacked only images, so the server saw the current pose repeated across history. Add `proprio_keys` (default empty) to `TemporalFrameStack` and `video_context_wrappers` so callers can stack e.g. `robot_state.ee_pose` + `grip` on the same buffer; each history frame then carries its own pose. Default stays empty so codecs that consume only the current proprio (e.g. DreamZero) are unaffected.
A `(T, H, W, 3)` frame stack is tens of MB raw and exceeds the ~2 MB websocket message cap of a Modal-fronted endpoint, so the obs never reaches the server. JPEG-encode 4D image stacks (q90 → ~0.9 MB for a 25-frame two-camera stack) into a compact marker the server decodes back to the same array; single frames stay raw. Also cuts upload latency on every request.
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ee6517f4a1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6829d2f502
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
2 participants
Summary
For policies that condition on a short window of history plus per-frame proprio:
TemporalStack(policy/wrappers.py) — replaces each named observation entry inkeyswith a(T, ...)temporal stack sampled at strided offsets, so the model sees a history window at the training cadence even though the harness only forwards an observation at re-query boundaries. It's modality-agnostic:keysare the cameras plus any per-frame proprio (e.g.robot_state.ee_pose,grip), each stacked step carrying its own value at that time.--policy.compress_images(default off), each image — a single(H, W, 3)frame or a(T, H, W, 3)stack — is JPEG-encoded per frame before send (q90) to stay under the ~2 MB websocket message cap of a Modal-fronted endpoint. Encode/decode live together inutils/serialization.py, and the marker decodes as a first-class wire type indeserialise, so every inference server restores the array transparently with no per-server change.video_context_wrappersdefault stacks images + proprio; DreamZero overrides to images-only.Test plan
Not yet verified end-to-end on hardware. Verified locally: the config instantiates, DreamZero resolves to images-only, the recorder handles stacked
(T, 7)/(T,)proprio without the earlier broadcast crash, and the offboard test suite passes — including a round-trip test that a JPEG-compressed image (single frame and stack) survivesserialise/deserialiseback to the original array and frame order.