Skip to content

feat(wam): add shared WAM backends and model/platform adapters - #250

Merged
qurakchin merged 50 commits into
RLinf:mainfrom
silence-breaker:feat/add-wam-backend
Nov 3, 2026
Merged

qurakchin merged 50 commits into
RLinf:mainfrom
silence-breaker:feat/add-wam-backend

Conversation

@silence-breaker

@silence-breaker silence-breaker commented Oct 10, 2026 •

Copy link
Copy Markdown
Collaborator

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/decode adapters: 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 --> Obs
Loading

For existing services, --am-endpoint replaces worker startup. Platform callbacks hand owned processes to the existing runtime for readiness, Dashboard status and cleanup.

Layer Responsibility
robots/<platform>/ Extract physical observations, define native control, connect clients, register tools and execute actions. RoboTwin's physical observation extractor is shared by LingBot and WAM.
am_client_base.py BaseAMClient shares inference dispatch/timeouts; BaseVLAClient and BaseWAMClient keep their respective request and result behavior.
am_facade_base.py BaseAMFacade shares prediction registration and RPC session setup. WAM adds paired adapters, metadata and reset hooks; VLA retains native routes.
am_contracts.py TypedDict declarations for physical observations/control and WAM metadata, requests and predictions. Runtime payloads remain dictionaries and NumPy arrays.
am_runtime.py Register common options, resolve AMConfig, and launch/connect workers using the existing daemon and RPC helpers.
components/<model>/ runtime_config.py resolves model deployment requirements; server.py loads a model or connects its native service; adapter/{__init__,encode,decode}.py maps platform/model interfaces.

WAM exposes wam.get_model_meta, wam.predict and wam.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.

Backend Supported platform Model-side conversion
Cosmos Policy LIBERO / LIBERO-Pro Two 256px RGB views, 9D state, and official inference producing 16 x 7 native actions
Cosmos3 LIBERO / LIBERO-Pro RLinf's SGLang adapter constructs requests and restores actions; its LIBERO action utility applies the gripper convention
Fast-WAM LIBERO / LIBERO-Pro Physical fields mapped to RLinf; RLinf owns preprocessing, inference and action restoration
Fast-WAM RoboTwin Three-camera composition, native joint targets and qpos14 actions

Fast-WAM LIBERO uses RLinf's get_model() with native RLinf YAML and predict_action_batch(); its model environment uses RLinf's supported FastWAM revision (45d8e14). Cosmos3 reuses Cosmos3SGLangAdapter, InferenceHTTPClient and 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-endpoint addresses SGLang, while --am-endpoint addresses 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 a MemoryManager. WAM defaults to a local memory directory scoped by platform/backend, which may be empty; --memory-dir can 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.5

Choose a planner model configured for your provider. --cuda-device 0 applies to local services, not the external worker. The bilingual CLI guide owns parameter and compatibility rules; add_action_model.rst covers architecture and model deployment, and platform guides cover task selection and execution.

Compatibility

The original --cuda-device remains the shared default: service-specific option > --cuda-device > inherited environment/service default. --am-cuda-device, --env-cuda-device and --sam3-cuda-device override their respective service where supported. Services may share a GPU.

VLA and WAM now share --am-backend, --am-endpoint, --am-checkpoint, --am-python, --am-root and --am-cuda-device. Previous option names remain hidden compatibility aliases and populate the same args.am_* fields:

Previous option Current option
--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-future enables 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 rlinf maps to --am-backend pi05; existing model-option aliases are retained. Select the backend before --help to 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 removed vla_client_base or vla_facade_base modules must instead import BaseVLAClient from am_client_base or BaseVLAFacade from am_facade_base; all in-repository callers are migrated.

Earlier RPent Fast-WAM LIBERO YAML needs migration to RLinf's model schema: execute_steps becomes num_action_chunks, inference options move to top-level model fields, and model/processor composition uses fastwam.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 base d903f38, head 5343c1e. Counts use git diff --numstat d903f38...5343c1e; each changed file appears once.

Area Files Added Deleted
Shared AM modules and removed VLA bases 6 +761 -100
Cosmos Policy worker and adapters 6 +411 -0
Cosmos3 RPC worker and RLinf adapters 6 +309 -0
Fast-WAM worker and adapters 6 +527 -0
LIBERO platform integration 10 +543 -130
RoboTwin platform integration 10 +399 -79
Other VLA callers and shared integration 25 +245 -194
Offline unit and RPC tests 25 +1568 -52
Opt-in GPU E2E tests and helpers 10 +226 -11
Bilingual documentation and READMEs 35 +1228 -310
Total 139 +6217 -876

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-future as 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.

Backend / platform Selection Result Pytest wall time
Cosmos Policy / LIBERO libero_spatial, task 0, seed 0 3 passed 61.43 s
Cosmos Policy / LIBERO-Pro libero_spatial_task, task 0, seed 0 3 passed 60.10 s
Fast-WAM / LIBERO libero_spatial, task 0, seed 0 3 passed 66.27 s
Fast-WAM / LIBERO-Pro libero_spatial_task, task 0, seed 0 3 passed 66.59 s
Fast-WAM / RoboTwin beat_block_hammer, demo_randomized, seed 100000 2 passed 189.03 s
Cosmos3 / LIBERO libero_10, task 0, seed 0 3 passed 95.38 s

Cosmos3's three checks were rerun at 726055a using the SGLang model identity. Fast-WAM's eight checks are from cff3e3c; 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, recorded finish and owned-process cleanup. They did not evaluate Qwen or memory-assisted reasoning. Workers were launched through start_wam_service(), and simulator/CLI tests borrowed their endpoints. Fast-WAM LIBERO uses RLinf 357f85b with its pinned FastWAM 45d8e14; RoboTwin uses official FastWAM 7faa711.

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 373e900 achieved 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 fa9310c completed 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.

Task Instruction Native success Actions Control time
0 Put alphabet soup and tomato sauce in the basket Yes 367 88.23 s
1 Put cream cheese and butter in the basket Yes 283 69.36 s
2 Turn on the stove and put the moka pot on it Yes 216 53.38 s
3 Put the black bowl in the bottom drawer and close it Yes 219 53.24 s
4 Put the white mug on the left plate and the yellow/white mug on the right plate Yes 224 53.75 s
5 Put the book in the back compartment of the caddy Yes 158 38.05 s
6 Put the white mug on the plate and pudding to its right Yes 222 53.98 s
7 Put alphabet soup and cream cheese in the basket Yes 234 57.39 s
8 Put both moka pots on the stove Yes 471 114.93 s
9 Put the yellow/white mug in the microwave and close it Yes 262 64.62 s

Conditions: nvidia/cosmos3-nano-libero10-sdg-export-S2_s43, official conversion with cosmos-framework c23e51f2; RLinf 357f85b, 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

  • Shared AM client/facade modules and deployment configuration, preserving VLA protocols.
  • Common --am-* options, legacy aliases and service-specific GPU overrides.
  • Paired Cosmos Policy / Cosmos3 / Fast-WAM adapters for the supported platforms listed above.
  • Shared boundary/configuration tests, model adapter tests and opt-in GPU integration tests.
  • Consolidate remaining VLA service launch/connection duplication while retaining native transports and lifecycle requirements.
  • Add further model/platform pairs, including RoboCasa, with compatible checkpoints and observation/control transforms.
  • Implement separate forecasting/planning operations and planner consumption of predicted observations/values.
  • Add Fast-WAM temporal action ensembling for supported checkpoints.
  • Make common runner client/session closure explicit at task teardown.
  • Connect opt-in WAM suites to the shared GPU-suite provisioning runner.
  • Add model-compatible WAM experience corpora and HF memory selection.

Copilot AI balanced review requested due to automatic review settings October 10, 2026 12:32

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

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.
@silence-breaker silence-breaker mentioned this pull request Oct 19, 2026
1 of 5 tasks
Comment thread rpent/robots/components/am_runtime.py Resolved
Comment thread rpent/robots/components/am_runtime.py Resolved
@silence-breaker
silence-breaker force-pushed the feat/add-wam-backend branch 2 times, most recently from b6d50f7 to 9d4d8b3 Compare October 29, 2026 01:53
Comment thread rpent/robots/robot_spec.py Resolved
Comment thread rpent/robots/robot_spec.py Resolved
@silence-breaker

Copy link
Copy Markdown
Collaborator Author

I have fixed the problems😀. Thanks for reviews. Are there any questions? @qurakchin

@silence-breaker

Copy link
Copy Markdown
Collaborator Author

Can you merge this PR now? or wait for @Sonorous281 's review?😀

@qurakchin
qurakchin merged commit 0fa63c5 into RLinf:main Nov 3, 2026
7 checks passed
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.

4 participants