Skip to content

feat: RTC and VLASH asynchronous inference - #2

Open
pppppppppper wants to merge 3 commits into
mainfrom
feat/async-inference-rtc-vlash
Open

pppppppppper wants to merge 3 commits into
mainfrom
feat/async-inference-rtc-vlash

Conversation

@pppppppppper

Copy link
Copy Markdown
Collaborator

feat: RTC and VLASH asynchronous inference

Adds the inference-side half of two published asynchronous-inference algorithms for
action-chunking VLA policies, in a new engine-layer feature domain
embodiinfer/engine/async_inference/.

Both address the same physical fact: in an asynchronous loop the chunk a policy
computes is not the chunk the robot is executing, because the robot keeps moving
during the forward pass.

  • RTC (Real-Time Chunking, arXiv 2506.07339) reconciles the actions — hold the
    actions guaranteed to execute and guide the remainder toward them.
  • VLASH (arXiv 2512.01031) reconciles the state — roll the robot state forward
    under the actions already issued and condition on the execution-time state.

Both are opt-in per request via one metadata["async"] block
(embodiinfer.async.v1); a request without it behaves byte-for-byte as before.

Motivation

AsyncEngine batches requests that happened to arrive together, but nothing in the
engine knows that a request is a re-plan for an episode still executing the previous
chunk
. Under asynchronous execution that chunk's unexecuted prefix is a commitment
the new chunk contradicts. The timing knowledge needed to reconcile them (how many
control steps of latency elapsed, which actions are still committed) lives on the
execution side, in a different process, behind the versioned HTTP/WirelessComm API —
so the algorithms are carried as a validated payload rather than shared objects.

docs/proposals/0002-gr00t-adapter.md already listed "RTC inpainting" as a
follow-up extension point.

Design

layer algorithm why there
PolicyService.step VLASH state roll-forward pure arithmetic on the request's own state mapping; needs no model, so every policy in the catalog inherits it
ActionDecoder / EngineCore RTC prefix guidance conditions the denoising trajectory, so it belongs to whoever owns the decode loop
  • rtc.py — the numeric core (prefix_weights, guidance_strength,
    rtc_guided_velocity, hard-prefix helpers, direction adapters).
  • vlash.py — roll_state_forward, roll_state_forward_projected,
    VlashStatePlan, apply_vlash_state_plan.
  • contracts.py / serving.py — the wire contract, validation, and the serving glue.
  • FlowDecoder.integrate(..., guidance=...) implements both conditioning modes;
    supports_rtc_guidance is a declared capability and a decoder that does not
    declare it rejects guidance rather than silently decoding without it.
  • The pi0.5 serving adapter keeps the per-session model-space committed prefix,
    because that is what the denoiser compares against and no client ever sees it.

Verification

Numeric fidelity is pinned to the reference implementations and to closed forms:

  • RTC prefix weights match Physical Intelligence's get_prefix_weights bit-for-bit
    (torch.equal), including the documented example
    (start=2, end=6, total=10 → 1 1 4/5 3/5 2/5 1/5 0 0 0 0).
  • The soft guidance correction is checked against an analytic Jacobian-vector
    product for an affine velocity field. This also pinned a real divergence: the
    LeRobot PyTorch port evaluates v_t before marking x_t as requiring grad, which
    collapses d x1/d x to the identity. This implementation follows PI's JAX
    original, and a test asserts the two differ so the choice cannot regress silently.
  • Both flow directions (t: 1→0 pi0.5/openpi, t: 0→1 GR00T/mock) are verified
    against the same closed form; the direction is derived from the schedule, not
    hardcoded.
  • rtc_guidance=None is torch.equal-identical to the pre-change engine output.
  • tests/fixtures/async_inference_golden.json is shared byte-identically with
    EmbodiRun, which asserts the same tables with its own numpy implementation — so a
    drift on either side of the HTTP boundary fails that side's suite.

Real-weight, two-process verification on 2×A100 with the π0.5 LIBERO checkpoint:
hard-prefix RTC pins the committed rows to the previously issued chunk
max|Δ| = 0.000e+00 through the whole chain (protocol → session cache → adapter →
engine → denoiser → postprocessor).

End-to-end LIBERO success rate (with EmbodiRun)

libero_spatial, 10 tasks × 5 trials, modeled delay Δ = 4 control steps
(a simulation parameter — the measured HTTP latency ≈ 200 ms is recorded separately
and never substituted for Δ). Each (task, seed) runs all modes with the same seed.

method inference checkpoint n success rate mean env steps
sync synchronous (upper bound) standard LIBERO 50 1.000 101.9
naive async async, no compensation standard LIBERO 50 0.620 325.2
RTC async + prefix guidance standard LIBERO 50 0.960 113.4
VLASH async + execution-time state VLASH delay16 fine-tune 50 0.780 245.6

Each method is evaluated on the checkpoint its design calls for. RTC recovers
essentially all of the synchronous performance (96% vs 100%, 113 vs 102 env steps);
naive loses a third of it; VLASH clearly beats naive (+16 points) but trails RTC.

Significance at n=50: RTC > VLASH is significant (48/50 vs 39/50, Fisher p ≈ 0.007);
VLASH > naive is not (39/50 vs 31/50, p ≈ 0.13) — the direction is right but 50
episodes do not resolve it.

Backward compatibility

Every new parameter is keyword-only and defaults to the current behaviour; the step
schema is untouched. A request with no async block, or rtc_guidance=None, takes
the original code path.

Notes for reviewers

  • Soft RTC guidance is a Jacobian-vector product per denoising step, so it runs
    eagerly and cannot be replayed from a captured CUDA graph. Hard-prefix mode needs
    no autograd. This is documented rather than hidden; capturing the backward pass is
    future work.
  • VLASH's published accuracy assumes offset fine-tuning. This PR implements the
    inference half only; the training augmentation is out of scope and the docs say so.
  • The pi0.5 committed-prefix cache is per-session process state, bounded by the
    session limit and dropped on reset.

Test plan

python -m pytest tests/test_async_rtc_vlash.py -q     # 68 tests
python -m pytest tests/test_serve_schema_aliases.py -q # 10 tests
python -m pytest tests/ -q
ruff check embodiinfer tests && ruff format --check embodiinfer tests

Full suite: 581 passed, 3 failed, 39 skipped. All 3 failures are pre-existing
environment failures, reproduced identically on a pristine git archive HEAD
tree (test_data_parallel.py::test_two_gpu_episode_affinity_with_per_replica_cuda_graphs,
test_lingbot_vla.py::test_lingbot_uses_explicit_backbone_for_weights_and_processor,
test_qwen25_vl_compile_cache.py::test_real_three_process_persistent_cache_admission).

The serving API was renamed from vvla to embodiinfer, and the two first-party
deployments ended up on opposite sides of it: this engine's examples and tests
send vvla.policy.*, while the EmbodiRun deployment runtime sends
embodiinfer.policy.*. The server only accepted the legacy spelling, so a real
EmbodiRun process could not open a session at all. Neither repository's own
tests could see this, because each only ever talks to itself.

Both spellings are now accepted on input for the session, step, structured-step
and RPC envelope schemas. Responses deliberately keep carrying the legacy
spelling, because a client asserting on it exists in-tree
(tests/test_http_client_example.py); flipping emission is a separate, deliberate
migration.

Found by running the two processes against each other.
Adds an engine-layer feature domain, embodiinfer/engine/async_inference/,
implementing the two published asynchronous-inference algorithms for
action-chunking VLA policies.

In an asynchronous loop the chunk a policy computes is not the chunk the robot
is executing, because the robot keeps moving during the forward pass. RTC
(arXiv 2506.07339) reconciles the actions: it holds the actions guaranteed to
execute and guides the remainder toward them. VLASH (arXiv 2512.01031)
reconciles the state: it rolls the robot state forward under the actions already
issued and conditions on the execution-time state.

Layering follows what each algorithm needs. VLASH is arithmetic on the request's
own state mapping, so it is applied once in PolicyService.step and every policy
in the catalog inherits it. RTC conditions the denoising trajectory, so it is a
declared decoder capability: ActionDecoder.supports_rtc_guidance, threaded
through EngineCore.execute as an optional per-batch argument. A decoder that does
not declare the capability rejects guidance rather than silently decoding without
it, because a chunk that contradicts the actions already committed to the robot is
a control hazard.

The pi0.5 serving adapter keeps the per-session committed prefix in model space,
which is what the denoiser compares against and what no client ever sees; the
request only has to carry the inference delay. Soft guidance is a Jacobian-vector
product per step, so it runs eagerly and is not graph-replayable; the hard-prefix
mode needs no autograd. Both are documented rather than hidden.

Numeric fidelity is pinned to Physical Intelligence's JAX original and to closed
forms, including an analytic Jacobian-vector product check that also pins a real
divergence in the LeRobot PyTorch port. tests/fixtures/async_inference_golden.json
is shared byte-identically with EmbodiRun and asserted by both, so the timing
arithmetic cannot drift across the process boundary unnoticed.

Verified end to end against EmbodiRun driving LIBERO with a real pi0.5
checkpoint: hard-prefix RTC pins the committed rows to the previously issued
chunk with max|delta| = 0 through protocol, session cache, adapter, engine,
denoiser and postprocessor.

See docs/proposals/0008-async-inference-rtc-vlash.md.
ZhouAo-ZA
ZhouAo-ZA previously approved these changes Oct 3, 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.

2 participants