diff --git a/configs/fastwam/robotwin_i2va.json b/configs/fastwam/robotwin_i2va.json index b80c77a5c..17bea6a97 100644 --- a/configs/fastwam/robotwin_i2va.json +++ b/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}" diff --git a/lightx2v/__init__.py b/lightx2v/__init__.py index a2250870d..1224dc922 100755 --- a/lightx2v/__init__.py +++ b/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__", diff --git a/lightx2v/models/networks/wan/fastwam_model.py b/lightx2v/models/networks/wan/fastwam_model.py index 4488b7f0d..c47a2c81d 100644 --- a/lightx2v/models/networks/wan/fastwam_model.py +++ b/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], diff --git a/lightx2v/models/runners/wan/fastwam_runner.py b/lightx2v/models/runners/wan/fastwam_runner.py index 7c1ddbdbf..655e5f1e8 100644 --- a/lightx2v/models/runners/wan/fastwam_runner.py +++ b/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 diff --git a/lightx2v_ros/src/robodojo/README.md b/lightx2v_ros/src/robodojo/README.md new file mode 100644 index 000000000..e7c0c79e5 --- /dev/null +++ b/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. diff --git a/lightx2v_ros/src/robodojo/package.xml b/lightx2v_ros/src/robodojo/package.xml new file mode 100644 index 000000000..a31537588 --- /dev/null +++ b/lightx2v_ros/src/robodojo/package.xml @@ -0,0 +1,13 @@ + + + + robodojo + 0.0.1 + RoboDojo/XPolicyLab adapter for LightX2V FastWAM inference. + user + Apache-2.0 + ament_python + + ament_python + + diff --git a/lightx2v_ros/src/robodojo/resource/robodojo b/lightx2v_ros/src/robodojo/resource/robodojo new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/lightx2v_ros/src/robodojo/resource/robodojo @@ -0,0 +1 @@ + diff --git a/lightx2v_ros/src/robodojo/robodojo/__init__.py b/lightx2v_ros/src/robodojo/robodojo/__init__.py new file mode 100644 index 000000000..8e252e18c --- /dev/null +++ b/lightx2v_ros/src/robodojo/robodojo/__init__.py @@ -0,0 +1,3 @@ +from .fastwam_adapter import FastWAMRoboDojoAdapter + +__all__ = ["FastWAMRoboDojoAdapter"] diff --git a/lightx2v_ros/src/robodojo/robodojo/fastwam_adapter.py b/lightx2v_ros/src/robodojo/robodojo/fastwam_adapter.py new file mode 100644 index 000000000..6528e47c3 --- /dev/null +++ b/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() diff --git a/lightx2v_ros/src/robodojo/setup.cfg b/lightx2v_ros/src/robodojo/setup.cfg new file mode 100644 index 000000000..24c342ea9 --- /dev/null +++ b/lightx2v_ros/src/robodojo/setup.cfg @@ -0,0 +1,4 @@ +[develop] +script_dir=$base/lib/robodojo +[install] +install_scripts=$base/lib/robodojo diff --git a/lightx2v_ros/src/robodojo/setup.py b/lightx2v_ros/src/robodojo/setup.py new file mode 100644 index 000000000..66c4489a2 --- /dev/null +++ b/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", +)