Skip to content

Add state handling to public VLA API - #50

Merged
LiangSu8899 merged 1 commit into
mainfrom
fix/public-vla-state-api
May 23, 2026
Merged

LiangSu8899 merged 1 commit into
mainfrom
fix/public-vla-state-api

Conversation

@LiangSu8899

Copy link
Copy Markdown
Member

Summary

  • attach state from VLAModel.predict() to observation dicts without overwriting explicit observation state
  • refresh prompt prefixes when set_prompt(..., state=...) frontends receive a changed state value
  • document the state contract for Pi0, Pi0.5, Pi0-FAST, GROOT N1.6, and GROOT N1.7

Validation

  • PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python3 -m pytest tests/test_load_model_use_fp8_kwarg.py -q (18 passed)
  • StableHLO: Pi0 RTX vs official PyTorch FP32 ref cosine 0.997490
  • StableHLO: Pi0.5 RTX vs official PyTorch FP32 ref cosine 0.999705
  • StableHLO: GROOT N1.6 RTX DiT/action path vs rtx_groot_ref.npz cosine 0.999248
  • StableHLO: Pi0-FAST JAX vs JAX bf16 ref worst segment cosine 0.993657
  • git diff --check

@LiangSu8899
LiangSu8899 merged commit dd35074 into main May 23, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant