Repository navigation
feat(wam): add shared WAM backends and model/platform adapters - #250
Merged
Merged
Conversation
Documentation build overview
28 files changed ·
|
Documentation build overview
28 files changed ·
|
Preserve RoboTwin memory hooks and WAM options, conditional E2E memory setup, and both LIBERO motion and WAM tests. Add a regression check that WAM skips HF memory preparation. Validation: 33 focused tests passed; 17 E2E cases collected; pre-commit passed. Full suite: 1504 passed, 5 skipped, 3 failed, 1 error. The same failures reproduce on untouched upstream/main in this environment.
1 of 5 tasks
qurakchin
reviewed
Oct 21, 2026
silence-breaker
force-pushed
the
feat/add-wam-backend
branch
2 times, most recently
from
October 29, 2026 01:53
b6d50f7 to
9d4d8b3
Compare
qurakchin
reviewed
Nov 2, 2026
Collaborator
Author
|
I have fixed the problems😀. Thanks for reviews. Are there any questions? @qurakchin |
qurakchin
approved these changes
Nov 3, 2026
wilburx813
approved these changes
Nov 3, 2026
Collaborator
Author
|
Can you merge this PR now? or wait for @Sonorous281 's review?😀 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add WAM backends with shared action-model infrastructure
Architecture
Add Cosmos Policy, Cosmos3 and Fast-WAM through reusable platform clients and paired model adapters. RPent owns agent orchestration, platform observation/control semantics and service wiring. RLinf owns supported model loading, preprocessing, inference and action restoration.
VLA integrations already reuse RPC bases, and Pi0.5 has an embodiment-specific encoder registry. WAM extends this reuse through physical-observation requests and model-side
encode/decodeadapters: compatible models share the platform client and execution tools. New WAMs normally add a model package and backend registry entry; new platform pairs still require matching checkpoints and transforms. VLA retains its native payloads, launchers and lifecycle extensions.flowchart TD Entry["CLI / Dashboard: --am-*"] --> Spec["Platform robot_spec.py"] Spec --> Config["am_runtime.select_am: AMConfig"] Config --> WAM["WAM: start_wam_service"] Config --> VLA["VLA: platform _spawn_vla_server"] WAM --> Spawn["AMConfig.spawn_worker: Python, cwd, GPU, logs"] VLA --> Spawn Spec --> Tool["Platform toolkit: wam_act"] Tool --> Obs["observation.py + wam_client.py"] Obs --> Client["BaseWAMClient in am_client_base.py"] Client <-->|"Existing RPent RPC"| Facade["BaseWAMFacade in am_facade_base.py"] Facade --> Encode["Model adapter/encode.py"] Encode --> Model["Model server: predict_native"] Model --> Decode["Model adapter/decode.py"] Decode --> Facade Client -->|"native actions"| Tool Tool --> Env["env_client.chunk_step / simulator"] Env --> ObsFor existing services,
--am-endpointreplaces worker startup. Platform callbacks hand owned processes to the existing runtime for readiness, Dashboard status and cleanup.robots/<platform>/am_client_base.pyBaseAMClientshares inference dispatch/timeouts;BaseVLAClientandBaseWAMClientkeep their respective request and result behavior.am_facade_base.pyBaseAMFacadeshares prediction registration and RPC session setup. WAM adds paired adapters, metadata and reset hooks; VLA retains native routes.am_contracts.pyTypedDictdeclarations for physical observations/control and WAM metadata, requests and predictions. Runtime payloads remain dictionaries and NumPy arrays.am_runtime.pyAMConfig, and launch/connect workers using the existing daemon and RPC helpers.components/<model>/runtime_config.pyresolves model deployment requirements;server.pyloads a model or connects its native service;adapter/{__init__,encode,decode}.pymaps platform/model interfaces.WAM exposes
wam.get_model_meta,wam.predictandwam.reset. The platform checks control compatibility once; the service validates incoming observations, and the client validates returned actions before execution. Optional timeout metadata aligns the client with the native service budget. Episode reset and session handling reuse existing RPC machinery.Fast-WAM LIBERO uses RLinf's
get_model()with native RLinf YAML andpredict_action_batch(); its model environment uses RLinf's supported FastWAM revision (45d8e14). Cosmos3 reusesCosmos3SGLangAdapter,InferenceHTTPClientand LIBERO action utilities. Cosmos Policy uses NVIDIA's official LIBERO implementation. Fast-WAM RoboTwin retains the official model path because RLinf's current wrapper lacks its three-camera/qpos14 support.Model packages own checkpoint requirements: Cosmos Policy and Fast-WAM load local weights; Cosmos3 identifies the model served by SGLang without requiring weights on the RPC host.
--cosmos3-sglang-endpointaddresses SGLang, while--am-endpointaddresses RPent RPC. SGLang owns its process and GPU allocation. Model dependencies live in separately provisioned environments, outside RPent's shared installation.Dashboard category selection reads the backend registry without validating deployment files. Task configuration and startup validate deployment requirements; fixed adapter layout/schema checks run at initialization. Per-prediction checks remain where malformed inputs or outputs can reach execution.
Both platforms expose
wam_act, with request-local instruction overrides and bounded chunk execution. Every toolkit requires aMemoryManager. WAM defaults to a local memory directory scoped by platform/backend, which may be empty;--memory-dircan override it. WAM currently excludes HF memory, Flash and exploration. Native simulator status determines task success. Optional future observations/values are carried through the WAM result, but are not used for planning.CLI Example
With the LIBERO simulator installed, a Cosmos Policy LIBERO worker already serving the endpoint, and planner credentials configured:
rpent --robot libero --suite libero_spatial --task 0 \ --am-backend cosmos-policy --am-endpoint http://127.0.0.1:8116 \ --cuda-device 0 --memory-profile local \ --planner api --model openai:gpt-5.5Choose a planner model configured for your provider.
--cuda-device 0applies to local services, not the external worker. The bilingual CLI guide owns parameter and compatibility rules;add_action_model.rstcovers architecture and model deployment, and platform guides cover task selection and execution.Compatibility
The original
--cuda-deviceremains the shared default: service-specific option >--cuda-device> inherited environment/service default.--am-cuda-device,--env-cuda-deviceand--sam3-cuda-deviceoverride their respective service where supported. Services may share a GPU.VLA and WAM now share
--am-backend,--am-endpoint,--am-checkpoint,--am-python,--am-rootand--am-cuda-device. Previous option names remain hidden compatibility aliases and populate the sameargs.am_*fields:--vla-backend,--wam-backend--am-backend--vla-endpoint,--wam-endpoint--am-endpoint--vla-model-path,--wam-checkpoint--am-checkpoint--pi05-python,--wam-python--am-python--wam-root--am-root--vla-cuda-device,--wam-cuda-device--am-cuda-device--predict-futureenables future-state/value outputs during local deployment, currently supported only by Cosmos Policy. Its former model-prefixed option names are removed. It does not enable planner consumption of these outputs; existing workers must be configured at deployment.RoboDojo's
--policy-backend rlinfmaps to--am-backend pi05; existing model-option aliases are retained. Select the backend before--helpto see its options. Explicit options for another model, or local deployment options combined with--am-endpoint, produce a configuration error.Existing VLA
vla.predict(obs, options)payloads and model-specific routes remain unchanged. VLA metadata interfaces are not replaced by the WAM protocol. Direct-action tools remain available. Python extensions importing the removedvla_client_baseorvla_facade_basemodules must instead importBaseVLAClientfromam_client_baseorBaseVLAFacadefromam_facade_base; all in-repository callers are migrated.Earlier RPent Fast-WAM LIBERO YAML needs migration to RLinf's model schema:
execute_stepsbecomesnum_action_chunks, inference options move to top-level model fields, and model/processor composition usesfastwam.config_dir/config_name/overrides. The worker rejects the old schema with a migration message. RoboTwin YAML and CLI aliases remain unchanged.Diff Breakdown
Compared with
RLinf/RPent:main, merge based903f38, head5343c1e. Counts usegit diff --numstat d903f38...5343c1e; each changed file appears once.Production code accounts for 69 files, +3195 / -503. This includes 761 lines in four shared AM modules, replacing two VLA base files totaling 100 lines, and 1247 lines for the three model packages. Cosmos3 adds 309 lines including module headers, configuration and metadata; its inference and numerical transforms reuse RLinf. Remaining production changes wire platforms, migrate VLA callers and consolidate deployment options.
Tests account for 35 files, +1794 / -63. Documentation accounts for 35 files, +1228 / -310. Together, tests and documentation account for approximately 49% of additions.
Tests
On 2026-10-27, after the checkpoint and Dashboard fixes, the full CPU unit suite passed: 1613 passed, 7 skipped, in 61.39 s, using Python 3.11. All pre-commit hooks passed. Skips cover unavailable optional RLinf imports and single-arm cases where dual-arm checks do not apply. This is one tested environment, not the full CI version matrix.
After removing duplicate Cosmos assertions and consolidating documentation at
2f3d88e, the affected adapter and Cosmos RPC tests passed: 17 passed, 3.57 s. Repository-wide pre-commit checks also passed; the full suite and GPU checks were not rerun for this cleanup.At
5343c1e, the AM and configuration tests passed (127 passed, 1.10 s) with--predict-futureas the sole future-prediction option, including worker forwarding and incompatible deployment options. Repository-wide pre-commit checks passed.Coverage includes configuration and legacy aliases, GPU option precedence, VLA dispatch, WAM boundaries, model transforms, platform execution, and real HTTP/socket sessions with reset/isolation. Regression cases cover missing local weights, model-identity forwarding and Dashboard display without deployment files.
GPU E2E
The existing platform component and CLI policy-chain tests passed across the following configurations. Fast-WAM was rerun on 2026-10-27 after adopting native RLinf configuration and moving RoboTwin invariant checks to startup; Cosmos Policy results are from 2026-10-26.
libero_spatial, task 0, seed 0libero_spatial_task, task 0, seed 0libero_spatial, task 0, seed 0libero_spatial_task, task 0, seed 0beat_block_hammer,demo_randomized, seed 100000libero_10, task 0, seed 0Cosmos3's three checks were rerun at
726055ausing the SGLang model identity. Fast-WAM's eight checks are fromcff3e3c; its inference path is unchanged and was not rerun for the checkpoint-validation fix. Timeout precedence is covered by unit tests.Hardware: A100-SXM4-40GB. These checks used real model weights and simulators, a scripted planner through the public CLI,
wam_act, state artifacts, recordedfinishand owned-process cleanup. They did not evaluate Qwen or memory-assisted reasoning. Workers were launched throughstart_wam_service(), and simulator/CLI tests borrowed their endpoints. Fast-WAM LIBERO uses RLinf357f85bwith its pinned FastWAM45d8e14; RoboTwin uses official FastWAM7faa711.LIBERO used a 32-action horizon; RoboTwin's CLI executed 48 actions within a 64-action horizon. Native task success was false in all bounded rollouts. Wall times exclude initial model loading and include simulator/test work; they are not benchmark success rates or inference latency measurements. Fresh installation of the complete GPU model/simulator environments remains unverified.
A historical policy-only Cosmos Policy evaluation at
373e900achieved 30/30 successes on LIBERO Spatial: all 10 tasks, initial states 0/1/2, horizon 220, 5 denoising steps and seed 1. It predates the current refactor.Cosmos3 Task Completion
On 2026-10-27, the Cosmos3 integration at
fa9310ccompleted 10/10 LIBERO-10 tasks, with zero runtime errors. This is policy-only execution through RPent's production LIBERO client, RLinf adapter and SGLang service. Native simulator success determines the outcome; execution stops at the successful action, including within a chunk.Conditions:
nvidia/cosmos3-nano-libero10-sdg-export-S2_s43, official conversion with cosmos-frameworkc23e51f2; RLinf357f85b, SGLang 0.5.19, Torch 2.13.0+cu129. Two A100-SXM4-40GB GPUs, TP=2, TORCH_SDPA, BF16, DiT CPU offload and VAE layerwise offload; a separate GPU for LIBERO. Each task used initial state 0, horizon 520, 30 denoising steps, policy seed 1 and 16-action predictions. No planner or memory guidance, future-state/value generation or text output.Mean control time was 64.69 s/episode, or 78.43 s including simulator startup/cleanup. Across 169 predictions, RPC latency was 3.324 s mean / 3.336 s p95; sampled peak model memory was 16,506 MiB per GPU. Timings exclude initial model loading and cold inference. Total output tokens were 0 for this continuous-action path. This covers one initial state per task and is not a full benchmark estimate.
Roadmap
--am-*options, legacy aliases and service-specific GPU overrides.