Add FastWAM RoboDojo evaluation adapter

configs/fastwam/robotwin_i2va.json

@@ -12,6 +12,7 @@
   "robot_state_dim": 14,
   "policy_profile": "robotwin",
   "normalize_mode": "z-score",
+  "target_video_length": 9,
   "binarize_gripper": false,
   "gripper_postprocess": false,
   "default_prompt": "A video recorded from a robot's point of view executing the following instruction: {task_prompt}"

lightx2v/__init__.py

@@ -4,7 +4,16 @@
 
 import lightx2v_platform.set_ai_device
 from lightx2v import common, models, utils
-from lightx2v.pipeline import LightX2VPipeline
+
+
+def __getattr__(name):
+    # Importing a model-specific runner (including ROS and RoboDojo adapters)
+    # should not eagerly require every optional dependency in pipeline.py.
+    if name == "LightX2VPipeline":
+        from lightx2v.pipeline import LightX2VPipeline
+
+        return LightX2VPipeline
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
 __all__ = [
     "__version__",

lightx2v/models/networks/wan/fastwam_model.py

@@ -42,6 +42,27 @@ def set_scheduler(self, scheduler):
         self.pre_infer.set_scheduler(scheduler)
         self.transformer_infer.set_scheduler(scheduler)
 
+    def _sequential_expert_offload_enabled(self):
+        return bool(self.config.get("sequential_aux_offload", False))
+
+    def prepare_video_phase(self):
+        """Keep only the video expert resident while building the KV cache."""
+        if not self._sequential_expert_offload_enabled():
+            return
+        self.transformer_weights.action.to_cpu()
+        self.transformer_weights.action_head.to_cpu()
+        torch.cuda.empty_cache()
+        self.transformer_weights.video.to_cuda()
+
+    def prepare_action_phase(self):
+        """Keep only the action expert resident during action denoising."""
+        if not self._sequential_expert_offload_enabled():
+            return
+        self.transformer_weights.video.to_cpu()
+        torch.cuda.empty_cache()
+        self.transformer_weights.action.to_cuda()
+        self.transformer_weights.action_head.to_cuda()
+
     def _load_ckpt(self, unified_dtype, sensitive_layer):
         adapter_path = self.config.get("adapter_model_path")
         if not adapter_path:
@@ -136,6 +157,7 @@ def prepare_action_inputs(
         context, context_mask = self._append_robot_state_to_context(context, context_mask, robot_state)
 
         video_pre, video_kv_cache = self._prepare_video_cache(first_frame_latents, context, context_mask)
+        self.prepare_action_phase()
         action_chunk_size = int(action_chunk_size)
         attention_mask = self.transformer_infer.build_mot_attention_mask(
             video_seq_len=video_pre.tokens.shape[0],

lightx2v/models/runners/wan/fastwam_runner.py

@@ -1,3 +1,4 @@
+import gc
 import json
 import os
 from collections import deque
@@ -158,6 +159,7 @@ def __init__(
             config=self.config,
             device=self.device,
         )
+        self.model.prepare_video_phase()
 
     @classmethod
     def from_config(cls, config):
@@ -198,25 +200,25 @@ def _load_normalizers(self):
         )
 
     def _load_text_encoder(self):
-        t5_path = self._find_model_file("models_t5_umt5-xxl-enc-bf16.pth")
-        tokenizer_path = self._find_model_dir("google/umt5-xxl")
+        t5_path = self._resolve_optional_file("t5_model_path", "models_t5_umt5-xxl-enc-bf16.pth")
+        tokenizer_path = self._resolve_optional_dir("tokenizer_path", "google/umt5-xxl")
         return T5EncoderModel(
             text_len=128,
             dtype=GET_DTYPE(),
             device=self.device,
             checkpoint_path=t5_path,
             tokenizer_path=str(tokenizer_path),
-            cpu_offload=False,
+            cpu_offload=bool(self.config.get("cpu_offload", False)),
         )
 
     def _load_vae(self):
-        vae_path = self._find_model_file("Wan2.2_VAE.pth")
+        vae_path = self._resolve_optional_file("vae_model_path", "Wan2.2_VAE.pth")
         return Wan2_2_VAE(
             vae_path=vae_path,
             device=self.device,
             dtype=GET_DTYPE(),
             vae_type="wan2.2",
-            cpu_offload=False,
+            cpu_offload=bool(self.config.get("cpu_offload", False)),
         )
 
     def _find_model_file(self, filename):
@@ -231,6 +233,28 @@ def _find_model_dir(self, dirname):
             return path
         raise FileNotFoundError(f"Cannot find {dirname} under {self.model_path}")
 
+    def _resolve_optional_file(self, config_key, default_filename):
+        value = self.config.get(config_key)
+        if not value:
+            return self._find_model_file(default_filename)
+        path = Path(str(value)).expanduser()
+        if not path.is_absolute():
+            raise ValueError(f"FastWAM requires absolute `{config_key}`, got: {path}")
+        if not path.is_file():
+            raise FileNotFoundError(str(path))
+        return str(path.resolve())
+
+    def _resolve_optional_dir(self, config_key, default_dirname):
+        value = self.config.get(config_key)
+        if not value:
+            return self._find_model_dir(default_dirname)
+        path = Path(str(value)).expanduser()
+        if not path.is_absolute():
+            raise ValueError(f"FastWAM requires absolute `{config_key}`, got: {path}")
+        if not path.is_dir():
+            raise FileNotFoundError(str(path))
+        return path.resolve()
+
     def reset(self):
         self.pending_actions.clear()
 
@@ -244,6 +268,7 @@ def next_action(self, images, state, task_description):
         return self.pending_actions.popleft()
 
     def predict_action_chunk(self, images, state, task_description, seed=None):
+        self.model.prepare_video_phase()
         image = self.build_image_tensor(images)
         first_frame_latents = self.encode_image_latents(image)
         context, context_mask = self.encode_prompt(self.default_prompt.format(task_prompt=task_description))
@@ -262,6 +287,12 @@ def predict_action_chunk(self, images, state, task_description, seed=None):
             seed=self.seed if seed is None else seed,
         )
         action = self.action_normalizer.backward(action).numpy()
+        if bool(self.config.get("sequential_aux_offload", False)):
+            # The result is on CPU; release per-plan CUDA tensors before the
+            # next replan when RoboDojo and inference share a small GPU.
+            del image, inputs, first_frame_latents, context, context_mask, robot_state
+            gc.collect()
+            torch.cuda.empty_cache()
         if self.gripper_postprocess:
             # LIBERO single-gripper channel: map [0,1] -> [-1,1], flip sign, optional binarize.
             action[..., -1] = action[..., -1] * 2 - 1

lightx2v_ros/src/robodojo/README.md

@@ -0,0 +1,30 @@
+# FastWAM on RoboDojo
+
+This package contains the observation/policy adapter used to evaluate
+LightX2V's native `FastWAMPolicy` through RoboDojo/XPolicyLab.
+
+The evaluator supplies three RGB observations:
+
+- `cam_head` -> `head_camera`
+- `cam_left_wrist` -> `left_camera`
+- `cam_right_wrist` -> `right_camera`
+
+The XPolicyLab policy wrapper should pack the robot state with its
+`pack_robot_state` helper, call `FastWAMRoboDojoAdapter.predict`, and unpack the
+returned 14-D absolute joint targets with `unpack_robot_state`. The released
+baseline-compatible settings are a 32-action chunk, 24 executed actions per
+plan, 10 action denoising steps, and z-score normalization.
+
+Required configuration paths are `model_path`, `checkpoint_path` (or
+`adapter_model_path`), and `dataset_stats_path`. When T5/tokenizer/VAE assets
+live outside `model_path`, pass absolute `t5_model_path`, `tokenizer_path`, and
+`vae_model_path` values. Set `sequential_aux_offload: true` when inference and
+Isaac Sim share a 24 GiB GPU.
+
+## Evaluation status
+
+The integration has completed end-to-end RoboDojo evaluation. Current results
+are close to the upstream FastWAM baseline: layout 0 succeeds, while layout 1
+still fails. Because the same layout-dependent behavior is suspected to come
+from the released official checkpoint, the current evidence does not attribute
+the layout 1 result to the LightX2V integration.

lightx2v_ros/src/robodojo/package.xml

@@ -0,0 +1,13 @@
+<?xml version="1.0"?>
+<?xml-model href="http://download.ros.org/schema/package_format3.xsd" schematypens="http://www.w3.org/2001/XMLSchema"?>
+<package format="3">
+  <name>robodojo</name>
+  <version>0.0.1</version>
+  <description>RoboDojo/XPolicyLab adapter for LightX2V FastWAM inference.</description>
+  <maintainer email="user@example.com">user</maintainer>
+  <license>Apache-2.0</license>
+  <buildtool_depend>ament_python</buildtool_depend>
+  <export>
+    <build_type>ament_python</build_type>
+  </export>
+</package>

lightx2v_ros/src/robodojo/resource/robodojo

@@ -0,0 +1 @@
+

lightx2v_ros/src/robodojo/robodojo/__init__.py

@@ -0,0 +1,3 @@
+from .fastwam_adapter import FastWAMRoboDojoAdapter
+
+__all__ = ["FastWAMRoboDojoAdapter"]

lightx2v_ros/src/robodojo/robodojo/fastwam_adapter.py

@@ -0,0 +1,84 @@
+"""Adapter between RoboDojo observations and LightX2V's FastWAM policy."""
+
+from pathlib import Path
+
+import numpy as np
+
+from lightx2v.models.runners.wan.fastwam_runner import FastWAMPolicy
+from lightx2v.utils.set_config import get_default_config
+
+
+def _rgb(image):
+    image = np.asarray(image)
+    if image.ndim != 3 or image.shape[-1] != 3:
+        raise ValueError(f"expected an HWC RGB image, got {image.shape}")
+    if image.dtype != np.uint8:
+        image = np.clip(image, 0, 255).astype(np.uint8)
+    return np.ascontiguousarray(image)
+
+
+class FastWAMRoboDojoAdapter:
+    """Small policy facade for XPolicyLab's RoboDojo server.
+
+    RoboDojo state packing is intentionally supplied by the caller because its
+    schema depends on ``env_cfg_type`` and ``action_type``. The adapter owns the
+    camera mapping and the FastWAM action-chunk/replan behavior.
+    """
+
+    camera_keys = {
+        "head_camera": "cam_head",
+        "left_camera": "cam_left_wrist",
+        "right_camera": "cam_right_wrist",
+    }
+
+    def __init__(self, config):
+        config = dict(config)
+        checkpoint = config.get("adapter_model_path") or config.get("checkpoint_path")
+        if not checkpoint:
+            raise ValueError("checkpoint_path (or adapter_model_path) is required")
+        stats = config.get("dataset_stats_path")
+        if not stats:
+            raise ValueError("dataset_stats_path is required")
+
+        runtime_config = get_default_config()
+        runtime_config.update(config)
+        runtime_config.update(
+            {
+                "model_cls": "fastwam",
+                "task": "i2va",
+                "adapter_model_path": str(Path(checkpoint).expanduser().resolve()),
+                "dataset_stats_path": str(Path(stats).expanduser().resolve()),
+                "policy_profile": "robotwin",
+                "normalize_mode": "z-score",
+                "action_dim": int(config.get("action_dim", 14)),
+                "robot_state_dim": int(config.get("robot_state_dim", 14)),
+                "action_chunk_size": int(config.get("action_chunk_size", 32)),
+                "actions_per_plan": int(config.get("actions_per_plan", 24)),
+                "action_infer_steps": int(config.get("action_infer_steps", 10)),
+                "action_sample_shift": float(config.get("action_sample_shift", 5.0)),
+                "default_prompt": config.get(
+                    "default_prompt",
+                    "A video recorded from a robot's point of view executing "
+                    "the following instruction: {task_prompt}",
+                ),
+            }
+        )
+        self.policy = FastWAMPolicy.from_config(runtime_config)
+
+    def predict(self, observation, packed_state, instruction):
+        vision = observation["vision"]
+        images = {
+            policy_key: _rgb(vision[robodojo_key]["color"])
+            for policy_key, robodojo_key in self.camera_keys.items()
+        }
+        return self.policy.predict_action_chunk(
+            images=images,
+            state=np.asarray(packed_state, dtype=np.float32),
+            task_description=str(instruction),
+        )[: self.policy.actions_per_plan]
+
+    def reset(self):
+        self.policy.reset()
+
+    def close(self):
+        self.policy.close()

lightx2v_ros/src/robodojo/setup.cfg

@@ -0,0 +1,4 @@
+[develop]
+script_dir=$base/lib/robodojo
+[install]
+install_scripts=$base/lib/robodojo

lightx2v_ros/src/robodojo/setup.py

@@ -0,0 +1,19 @@
+from setuptools import find_packages, setup
+
+package_name = "robodojo"
+
+setup(
+    name=package_name,
+    version="0.0.1",
+    packages=find_packages(),
+    data_files=[
+        ("share/ament_index/resource_index/packages", ["resource/" + package_name]),
+        ("share/" + package_name, ["package.xml"]),
+    ],
+    install_requires=["setuptools"],
+    zip_safe=True,
+    maintainer="user",
+    maintainer_email="user@example.com",
+    description="RoboDojo/XPolicyLab adapter for LightX2V FastWAM inference.",
+    license="Apache-2.0",
+)