Add Wall-X serving and Turtle2 TCP WebSocket bridge

wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper.py

@@ -19,10 +19,12 @@ from wall_x._vendor.x2robot_utils.grounding import (
 )
 
 from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig
-from wall_x._vendor.harrix.utils.ckpt_load import reshape_compatible_state_dict
+from wall_x._vendor.harrix.utils.ckpt_load import (
+    reshape_compatible_state_dict,
+    resolve_lora_scale,
+)
 from wall_x._vendor.harrix.utils.train_config import (
     resolve_camera_label,
-    resolve_max_length,
     resolve_state_bins,
     resolve_use_state_string_representation,
 )
@@ -165,8 +167,17 @@ class WallxModelWrapper:
                 "The weights is fused, skipping conversion.",
             )
 
+        lora_scale = resolve_lora_scale(
+            checkpoint_path,
+            self.config.train_config,
+            state_dict,
+            log_fn=self.logger.warning,
+        )
         state_dict = reshape_compatible_state_dict(
-            state_dict, self.model.state_dict(), log_fn=self.logger.info
+            state_dict,
+            self.model.state_dict(),
+            log_fn=self.logger.info,
+            lora_scale=lora_scale,
         )
         msg = self.model.load_state_dict(state_dict, strict=False)
         self.model.set_normalizer(
@@ -319,9 +330,11 @@ class WallxModelWrapper:
                 images=image_inputs,
                 videos=None,
                 padding=True,
-                truncation=True,
+                # Keep every image placeholder during online inference. The
+                # training text limit can cut them for three large images.
+                truncation=False,
                 return_tensors="pt",
-                max_length=resolve_max_length(self.config.train_config),
+                max_length=None,
                 pad_to_128_multiple=False,
                 pad_prefix_to_same_length=pad_prefix,
                 norm_state=(

wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper_rtc.py

@@ -0,0 +1,802 @@
+import os
+import torch
+import copy
+from safetensors.torch import load_file
+import numpy as np
+from PIL import Image
+from qwen_vl_utils.vision_process import smart_resize
+from transformers import BatchFeature
+
+from wall_x.trainer.trainer_utils import load_wallx_processors
+
+from wall_x._vendor.x2robot_utils.text_templates import (
+    preprocesser_call,
+    get_prologue_with_embodied_information,
+)
+from wall_x._vendor.x2robot_utils.grounding import (
+    reverse_grounding_points,
+    extract_grounding_points,
+)
+
+from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig
+from wall_x._vendor.harrix.utils.ckpt_load import (
+    reshape_compatible_state_dict,
+    resolve_lora_scale,
+)
+from wall_x._vendor.harrix.utils.train_config import (
+    resolve_camera_label,
+    resolve_state_bins,
+    resolve_use_state_string_representation,
+)
+from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger
+from wall_x.utils.timers import timer, ScopeTimer
+
+
+ENABLE_FAST_PREPROCESS = os.getenv("ENABLE_FAST_PREPROCESS", "False").lower() == "true"
+
+
+def move_to_cuda(obj, device="cuda"):
+    if isinstance(obj, torch.Tensor):
+        return obj.to(device)
+    elif isinstance(obj, (dict, BatchFeature)):
+        return {k: move_to_cuda(v, device) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [move_to_cuda(v, device) for v in obj]
+    elif isinstance(obj, tuple):
+        return tuple(move_to_cuda(v, device) for v in obj)
+    else:
+        return obj
+
+
+class WallxModelWrapper:
+    def __init__(self, config: InferConfig, rtc_config=None):
+        from wall_x._vendor.harrix.serving.rtc_wallx import (
+            WallXRTCConfig,
+            WallXRTCProcessor,
+        )
+
+        self.config = config
+        self.rtc_config = rtc_config or WallXRTCConfig()
+        self.rtc_processor = WallXRTCProcessor(self.rtc_config)
+        self.logger = InferLogger.get_model_logger("WallxModelWrapper")
+        self.norm_key = self.config.norm_key
+        self._register_normalizers()
+        self.logger.info(f"normalizers {self.norm_key} registered")
+        self._load_processor()
+        self._load_model()
+        self.load_ckpt()
+        for parameter in self.model.parameters():
+            parameter.requires_grad_(False)
+        self.logger.info(
+            "model %s loaded; parameters frozen for RTC input-VJP inference",
+            self.config.checkpoint_path,
+        )
+
+        self.norm_key = self.config.norm_key
+        self._register_normalizers()
+        self.logger.info(f"normalizers {self.norm_key} registered")
+
+        # Initialize robot_type_id (for v3.1 delta tokenizer)
+        self.robot_type_id = None
+        if self.tokenizer_mixin is not None:
+            robot_type = getattr(self.config, "robot_type", None)
+            if robot_type:
+                self.tokenizer_mixin.init_inference(robot_type=robot_type)
+                if hasattr(self.tokenizer_mixin, "robot_type_id"):
+                    self.robot_type_id = self.tokenizer_mixin.robot_type_id
+                    if self.robot_type_id is not None:
+                        self.logger.info(
+                            f"Initialized robot_type_id: {self.robot_type_id} from robot_type: {robot_type}"
+                        )
+
+        self.role_start_symbol = "<|im_start|>"
+        self.role_end_symbol = "<|im_end|>"
+        self.vision_start_symbol = "<|vision_start|>"
+        self.vision_end_symbol = "<|vision_end|>"
+        self.image_pad_symbol = "<|image_pad|>"
+        self.propri_symbol = "<|propri|>"
+        self.action_symbol = "<|action|>"
+
+        self.cam_names = self.config.cam_names
+
+    def _camera_name_mapping(self):
+        return self.config.train_config.get("data", {}).get("camera_name_mapping")
+
+    def _load_processor(self):
+        # Load tokenizer on model_device for inference
+        device = getattr(self.config, "model_device", "cuda")
+        processors_dict = load_wallx_processors(self.config.train_config, device=device)
+        self.processor = processors_dict["processor"]
+        self.action_mapper = processors_dict["action_mapper"]
+        self.tokenizer_mixin = processors_dict.get("tokenizer_mixin")
+
+    def _load_model(self):
+        from wall_x.trainer.adapters import resolve_adapter
+        from wall_x.model.qact.qwen2_5.modeling_qwen2_5_vl_act_rtc import (
+            Qwen2_5_VLMoEForAction,
+        )
+
+        model_type = self.config.train_config["model_type"]
+        if model_type != "qwen2_5":
+            raise ValueError(
+                f"RTC serving currently supports model_type='qwen2_5', got {model_type!r}"
+            )
+        adapter_cls = resolve_adapter(model_type)
+        ModelClass = Qwen2_5_VLMoEForAction
+        self.ModelClass = ModelClass
+
+        self.logger.info(
+            "initializing RTC model: %s (%s)", model_type, ModelClass.__name__
+        )
+        self.model = ModelClass(
+            self.config.model_config,
+            self.processor,
+            self.tokenizer_mixin,
+        )
+        adapter_cls.log_attention_implementation(self.logger, self.model)
+        self.model_type = model_type
+
+        self.logger.info("resizing model token embeddings")
+        self.model.resize_token_embeddings(len(self.processor.tokenizer))
+        self.logger.info("token embedding resize done")
+        self.logger.info("casting selected params to bfloat16")
+        self.model.to_bfloat16_for_selected_params()
+        self.logger.info("bfloat16 cast done")
+
+    def load_ckpt(self, checkpoint_path: str = None):
+        if checkpoint_path is None:
+            checkpoint_path = self.config.checkpoint_path
+
+        if os.path.exists(os.path.join(checkpoint_path, "global_step.pth")):
+            global_step = torch.load(os.path.join(checkpoint_path, "global_step.pth"))[
+                "global_step"
+            ]
+            self.logger.info(f"checkpoint global_step: {global_step}")
+
+        fsdp_ckpt = os.path.join(checkpoint_path, "pytorch_model_fsdp.bin")
+        safetensor_ckpt = os.path.join(checkpoint_path, "model.safetensors")
+        if os.path.exists(fsdp_ckpt):
+            self.logger.info(f"Loading FSDP checkpoint: {fsdp_ckpt}")
+            state_dict = torch.load(fsdp_ckpt, map_location="cpu")
+            # Unwrap outer dict (e.g. {'state_dict': {...}})
+            if isinstance(state_dict, dict) and "state_dict" in state_dict:
+                self.logger.info(
+                    "Using nested state_dict inside pytorch_model_fsdp.bin"
+                )
+                state_dict = state_dict["state_dict"]
+        elif os.path.exists(safetensor_ckpt):
+            self.logger.info(f"loading safetensors checkpoint: {safetensor_ckpt}")
+            state_dict = load_file(safetensor_ckpt, device="cpu")
+        else:
+            raise FileNotFoundError(
+                f"ERROR: No checkpoint found under {checkpoint_path}. "
+                "Expecting either pytorch_model_fsdp.bin or model.safetensors."
+            )
+
+        # Qwen models need fused weight conversion
+        if not self.ModelClass.is_fused(state_dict):
+            self.logger.info(
+                "Converting non-fused weights to fused format...",
+            )
+            state_dict = self.ModelClass.convert_to_fused(state_dict)
+        else:
+            self.logger.info(
+                "The weights is fused, skipping conversion.",
+            )
+
+        lora_scale = resolve_lora_scale(
+            checkpoint_path,
+            self.config.train_config,
+            state_dict,
+            log_fn=self.logger.warning,
+        )
+        state_dict = reshape_compatible_state_dict(
+            state_dict,
+            self.model.state_dict(),
+            log_fn=self.logger.info,
+            lora_scale=lora_scale,
+        )
+        msg = self.model.load_state_dict(state_dict, strict=False)
+        self.model.set_normalizer(
+            copy.deepcopy(self.normalizer_action),
+            copy.deepcopy(self.normalizer_propri),
+        )
+        self.logger.info(f"load_state_dict result: {msg}")
+        self.model.eval()
+        self.model.to(self.config.model_device)
+        self.model.to_bfloat16_for_selected_params()
+
+        if hasattr(self.model, "load_optimized_weights"):
+            self.model.load_optimized_weights(state_dict)
+
+    def _register_normalizers(self):
+        from wall_x._vendor.harrix.utils.normalizer import build_normalizers
+
+        self.normalizer_action, self.normalizer_propri, resolved = build_normalizers(
+            self.config.checkpoint_path,
+            self.config.train_config,
+            self.config.norm_key,
+        )
+        self.norm_key = resolved
+        self.config.norm_key = resolved
+        self._log_norm_debug(self.norm_key)
+
+    def _log_norm_debug(self, norm_key):
+        if not hasattr(self, "normalizer_action") or not hasattr(
+            self, "normalizer_propri"
+        ):
+            self.logger.warning("[NormDebug] normalizer is not initialized yet")
+            return
+
+        if (
+            norm_key in self.normalizer_action.min
+            and norm_key in self.normalizer_propri.min
+        ):
+            self.logger.debug(
+                "[NormDebug] norm_key=%s action(min/delta) shape=%s/%s",
+                norm_key,
+                tuple(self.normalizer_action.min[norm_key].shape),
+                tuple(self.normalizer_action.delta[norm_key].shape),
+            )
+            self.logger.debug(
+                "[NormDebug] action min(all)=%s delta(all)=%s",
+                self.normalizer_action.min[norm_key].detach().cpu().tolist(),
+                self.normalizer_action.delta[norm_key].detach().cpu().tolist(),
+            )
+            self.logger.debug(
+                "[NormDebug] propri(min/delta) shape=%s/%s",
+                tuple(self.normalizer_propri.min[norm_key].shape),
+                tuple(self.normalizer_propri.delta[norm_key].shape),
+            )
+            self.logger.debug(
+                "[NormDebug] propri min(all)=%s delta(all)=%s",
+                self.normalizer_propri.min[norm_key].detach().cpu().tolist(),
+                self.normalizer_propri.delta[norm_key].detach().cpu().tolist(),
+            )
+        else:
+            self.logger.warning(
+                "[NormDebug] norm_key=%s not found in normalizer", norm_key
+            )
+
+    @timer
+    def construct_model_input(
+        self, observation, prefix_text, postfix_text, pad_prefix=False
+    ):
+        batch_size = len(observation)
+        dataset_names = [self.norm_key] * batch_size
+        self.logger.debug("[NormDebug] dataset_names=%s", dataset_names)
+
+        additional_inputs = {}
+
+        # -------- proprioception / masks (batch) --------
+        agent_pos_list = []
+        agent_pos_mask_list = []
+        dof_mask_list = []
+        for obs in observation:
+            if "robot_state_action_data" in obs:
+                robot_state_action_data = obs["robot_state_action_data"]
+
+                agent_pos = torch.from_numpy(robot_state_action_data.agent_pos)
+                # Normalize to [1, T, D] for batch cat
+                if agent_pos.dim() == 2:
+                    agent_pos = agent_pos.unsqueeze(0)
+                agent_pos_list.append(agent_pos)
+
+                agent_pos_mask = torch.from_numpy(
+                    robot_state_action_data.agent_pos_mask
+                )
+                if agent_pos_mask.dim() == 2:
+                    agent_pos_mask = agent_pos_mask.unsqueeze(0)
+                agent_pos_mask_list.append(agent_pos_mask)
+
+                dof_mask = torch.from_numpy(robot_state_action_data.dof_mask)
+                if dof_mask.dim() == 1:
+                    dof_mask = dof_mask.unsqueeze(0)
+                dof_mask_list.append(dof_mask)
+
+        # cat: [B, T, D] / [B, ...]
+        if len(agent_pos_list) > 0:
+            agent_pos = torch.cat(agent_pos_list, dim=0)
+            agent_pos_mask = torch.cat(agent_pos_mask_list, dim=0)
+            dof_mask = torch.cat(dof_mask_list, dim=0)
+
+            if self.normalizer_propri is not None:
+                agent_pos = self.normalizer_propri.normalize_data(
+                    agent_pos, dataset_names
+                )
+
+            additional_inputs["proprioception"] = agent_pos.detach()
+            additional_inputs["agent_pos_mask"] = agent_pos_mask
+            additional_inputs["dof_mask"] = dof_mask
+
+        # -------- images (flattened, in placeholder scan order) --------
+        with ScopeTimer("resize_images"):
+            # TODO[KC]: optimize this using batch processing
+            image_inputs = []
+            all_image_sizes = []
+            if ENABLE_FAST_PREPROCESS:
+                for obs in observation:
+                    current_image_inputs = self._resize_images_fast(obs)
+                    image_inputs.extend(current_image_inputs)
+                    # tensor shape is (H, W, C), convert to (W, H) to match PIL.size format
+                    all_image_sizes.extend(
+                        [
+                            (image_i.shape[1], image_i.shape[0])
+                            for image_i in current_image_inputs
+                        ]
+                    )
+            else:
+                for obs in observation:
+                    current_image_inputs = self._resize_images(obs)
+                    image_inputs.extend(current_image_inputs)
+                    # tensor shape is (H, W, C), convert to (W, H) to match PIL.size format
+                    all_image_sizes.extend(
+                        [
+                            (image_i.shape[1], image_i.shape[0])
+                            for image_i in current_image_inputs
+                        ]
+                    )
+
+        additional_inputs["image_size"] = all_image_sizes
+
+        with ScopeTimer("preprocesser_call"):
+            inputs = preprocesser_call(
+                processor=self.model.processor,
+                prefix_text=prefix_text,
+                postfix_text=postfix_text,
+                images=image_inputs,
+                videos=None,
+                padding=True,
+                # Keep every image placeholder during online inference. The
+                # training text limit can cut them for three large images.
+                truncation=False,
+                return_tensors="pt",
+                max_length=None,
+                pad_to_128_multiple=False,
+                pad_prefix_to_same_length=pad_prefix,
+                norm_state=(
+                    additional_inputs["proprioception"]
+                    if "proprioception" in additional_inputs
+                    and resolve_use_state_string_representation(
+                        self.config.train_config
+                    )
+                    else None
+                ),
+                agent_pos_mask=(
+                    additional_inputs["agent_pos_mask"]
+                    if "agent_pos_mask" in additional_inputs
+                    else None
+                ),
+                state_augmentation_prob=0.0,
+                state_drop_prob=0.0,
+                state_augmentation_ratio=0.0,
+                state_bins=resolve_state_bins(self.config.train_config),
+                inference_mode=True,
+            )
+
+        with ScopeTimer("convert_action_token_id and post "):
+            action_token_id = self.model.processor.tokenizer.convert_tokens_to_ids(
+                "<|action|>"
+            )
+            flow_action_mask = inputs["input_ids"] == action_token_id
+            additional_inputs["moe_token_types"] = flow_action_mask
+            additional_inputs["dataset_names"] = dataset_names
+
+            inputs.update(additional_inputs)
+            inputs = move_to_cuda(inputs, self.config.model_device)
+
+        return inputs
+
+    def get_text_for_dllm_action(self, instruction):
+        if (
+            self.config.train_config["data"].get("use_embodied_system_prompt_ratio", 0)
+            > 0
+        ):
+            if self.norm_key != "x2_normal" and self.norm_key != "ex_normal":
+                self.config.robot_id = 0
+            cam_name_mapping = {cam_name: cam_name for cam_name in self.cam_names}
+            prologue = get_prologue_with_embodied_information(
+                dataset_name=self.norm_key,
+                cam_mapping=cam_name_mapping,
+                robot_id=self.config.robot_id,
+                uid="",
+                config=self.config.data_config,
+            )
+        else:
+            prologue = f"{self.role_start_symbol}system\nYou are a helpful assistant.{self.role_end_symbol}\n"
+        user_request = f"{self.role_start_symbol}user\nObservation:"
+        camera_name_mapping = self._camera_name_mapping()
+        for cam_name in self.cam_names:
+            user_request += (
+                f" {resolve_camera_label(cam_name, camera_name_mapping)}:"
+                f" {self.vision_start_symbol}{self.image_pad_symbol}{self.vision_end_symbol}"
+            )
+        user_request += "\nInstruction:"
+        text_prompt = f"\nPredict the next action in robot action.\nProprioception: {self.propri_symbol}\n"
+        user_message = (
+            f"{user_request} {instruction}{text_prompt}{self.role_end_symbol}\n"
+        )
+        placeholder_seq = self.tokenizer_mixin.get_placeholder_for_dllm()
+        ar_token = "".join(placeholder_seq) + "<|im_end|>\n"
+        assistant_message = f"{self.role_start_symbol}assistant\n{ar_token}"
+        flow_action = f"{self.action_symbol * self.config.action_horizon}"
+
+        prefix_text = prologue + user_message + assistant_message
+        postfix_text = flow_action
+
+        return prefix_text, postfix_text
+
+    @timer
+    def get_text_for_action(self, instruction):
+        if (
+            self.config.train_config["data"].get("use_embodied_system_prompt_ratio", 0)
+            > 0
+        ):
+            if self.norm_key not in ["x2_normal", "ex_normal"]:
+                self.config.robot_id = 0
+            cam_name_mapping = {cam_name: cam_name for cam_name in self.cam_names}
+            prologue = get_prologue_with_embodied_information(
+                dataset_name=self.norm_key,
+                cam_mapping=cam_name_mapping,
+                robot_id=self.config.robot_id,
+                uid="",
+                config=self.config.data_config,
+            )
+        else:
+            prologue = f"{self.role_start_symbol}system\nYou are a helpful assistant.{self.role_end_symbol}\n"
+        user_request = f"{self.role_start_symbol}user\nObservation:"
+        camera_name_mapping = self._camera_name_mapping()
+        for cam_name in self.cam_names:
+            user_request += (
+                f" {resolve_camera_label(cam_name, camera_name_mapping)}:"
+                f" {self.vision_start_symbol}{self.image_pad_symbol}{self.vision_end_symbol}"
+            )
+        user_request += "\nInstruction:"
+        text_prompt = f"\nPredict the next action in robot action.\nProprioception: {self.propri_symbol}\n"
+        user_message = (
+            f"{user_request} {instruction}{text_prompt}{self.role_end_symbol}\n"
+        )
+        assistant_message = f"{self.role_start_symbol}assistant\n"
+        flow_action = f"{self.action_symbol * self.config.action_horizon}"
+
+        prefix_text = prologue + user_message + assistant_message
+        postfix_text = flow_action
+
+        return prefix_text, postfix_text
+
+    @timer
+    def get_text_for_subtask_generation(self, instruction):
+        prologue = f"{self.role_start_symbol}system\nYou are a helpful assistant.{self.role_end_symbol}\n"
+        user_request = f"{self.role_start_symbol}user\nObservation:"
+        camera_name_mapping = self._camera_name_mapping()
+        for cam_name in self.cam_names:
+            user_request += (
+                f" {resolve_camera_label(cam_name, camera_name_mapping)}:"
+                f" {self.vision_start_symbol}{self.image_pad_symbol}{self.vision_end_symbol}"
+            )
+        user_request += "\nInstruction:"
+        text_prompt = "\nPredict the next action in language.\n"
+        user_message = (
+            f"{user_request} {instruction}{text_prompt}{self.role_end_symbol}\n"
+        )
+        assistant_message = f"{self.role_start_symbol}assistant\n"
+
+        prefix_text = prologue + user_message + assistant_message
+        postfix_text = ""
+
+        return prefix_text, postfix_text
+
+    @timer
+    def _resize_images(self, observation):
+        image_inputs = []
+        for key in self.cam_names:
+            if key not in observation:
+                continue
+            current_obs = observation[key]
+            if isinstance(current_obs, np.ndarray):
+                img_pil = Image.fromarray(current_obs)
+            elif isinstance(current_obs, Image.Image):
+                img_pil = current_obs
+            else:
+                raise ValueError(f"Unsupported image type: {type(current_obs)}")
+            orig_width, orig_height = img_pil.size
+
+            target_size = self.config.data_config.resolution.get(key, -1)
+            if target_size != -1:
+                # Aspect-ratio-preserving resize
+                if orig_width > orig_height:  # landscape
+                    new_width = target_size
+                    new_height = int(target_size * orig_height / orig_width)
+                else:  # portrait
+                    new_height = target_size
+                    new_width = int(target_size * orig_width / orig_height)
+                img_pil = img_pil.resize((new_width, new_height))
+
+            # Apply smart resize (Qwen logic)
+            current_width, current_height = img_pil.size
+            resized_height, resized_width = smart_resize(
+                current_height,
+                current_width,
+                factor=self.config.data_config.image_factor,
+                min_pixels=self.config.data_config.min_pixels,
+                max_pixels=self.config.data_config.max_pixels,
+            )
+            resized_img = img_pil.resize((resized_width, resized_height))
+            resized_img = torch.from_numpy(np.array(resized_img)).to(
+                self.config.model_device
+            )
+            image_inputs.append(resized_img)
+
+        return image_inputs
+
+    def _resize_images_fast(self, observation):
+        import cv2
+
+        image_inputs = []
+        for key in self.cam_names:
+            if key not in observation:
+                continue
+            current_obs = observation[key]
+            orig_height, orig_width, _ = current_obs.shape
+
+            target_size = self.config.data_config.resolution.get(key, -1)
+            current_width, current_height = orig_width, orig_height
+            if target_size != -1:
+                # Aspect-ratio-preserving resize
+                if orig_width > orig_height:  # landscape
+                    new_width = target_size
+                    new_height = int(target_size * orig_height / orig_width)
+                else:  # portrait
+                    new_height = target_size
+                    new_width = int(target_size * orig_width / orig_height)
+                current_width = new_width
+                current_height = new_height
+
+            # Apply smart resize (Qwen logic)
+            resized_height, resized_width = smart_resize(
+                current_height,
+                current_width,
+                factor=self.config.data_config.image_factor,  # FIXME
+                min_pixels=self.config.data_config.min_pixels,  # FIXME
+                max_pixels=self.config.data_config.max_pixels,  # FIXME
+            )
+
+            resized_img = cv2.resize(
+                current_obs,
+                (resized_width, resized_height),
+                interpolation=cv2.INTER_CUBIC,
+            )
+            resized_img = torch.from_numpy(resized_img).to(self.config.model_device)
+            image_inputs.append(resized_img)
+
+        return image_inputs
+
+    def infer_flow_action(
+        self,
+        observation,
+        instruction,
+        *,
+        prev_chunk_left_over=None,
+        inference_delay=0,
+        execution_horizon=None,
+    ):
+        self.logger.info("generating flow action")
+        self.logger.info(f"flow action instruction: {instruction}")
+
+        prefix_text, postfix_text = self.get_text_for_action(instruction)
+        model_input = self.construct_model_input(
+            [observation], [prefix_text], [postfix_text]
+        )
+
+        padding = (
+            torch.zeros_like(
+                self.normalizer_action.delta[model_input["dataset_names"][0]]
+            )
+            .unsqueeze(0)
+            .to("cpu")
+        )
+        padding_action = self.normalizer_action.normalize_data(
+            padding, model_input["dataset_names"]
+        ).to(model_input["input_ids"].device)
+
+        self.logger.info(
+            "generate_flow_action start (horizon=%s, flow_steps=%s, device=%s)",
+            self.config.action_horizon,
+            self.config.num_inference_timesteps,
+            self.config.model_device,
+        )
+        if torch.cuda.is_available():
+            try:
+                free_b, total_b = torch.cuda.mem_get_info(
+                    torch.device(self.config.model_device)
+                )
+                self.logger.info(
+                    "CUDA mem before flow: free=%.2f GiB / total=%.2f GiB",
+                    free_b / (1024**3),
+                    total_b / (1024**3),
+                )
+            except Exception as e:
+                self.logger.warning("CUDA mem_get_info failed: %s", e)
+
+        with ScopeTimer("generate_flow_action"):
+            model_output = self.model.generate_flow_action(
+                action_horizon=self.config.action_horizon,
+                action_dim=self.config.action_dim,
+                num_inference_timesteps=self.config.num_inference_timesteps,
+                padding_action=padding_action,
+                rtc_processor=self.rtc_processor,
+                prev_chunk_left_over=prev_chunk_left_over,
+                inference_delay=inference_delay,
+                execution_horizon=execution_horizon,
+                **model_input,
+            )
+
+        self.logger.info("flow action generation done")
+
+        model_output["robot_state_action_data"] = observation["robot_state_action_data"]
+        model_output["robot_state_action_data"].save_action_data(
+            model_output["predict_action"]
+        )
+        self.logger.info("saved flow action to robot_state_action_data")
+        return model_output
+
+    def infer_flow_action_batch(self, observations, instructions):
+        """
+        Batch flow action inference:
+        - observations: List[Dict], same format as single inference
+        - instructions: List[str], aligned with observations
+        Returns List[model_output] of length batch size
+        """
+        assert len(observations) == len(
+            instructions
+        ), "observations and instructions must have the same length"
+        batch_size = len(observations)
+
+        prefix_list = []
+        postfix_list = []
+        for ins in instructions:
+            prefix_text, postfix_text = self.get_text_for_action(ins)
+            prefix_list.append(prefix_text)
+            postfix_list.append(postfix_text)
+
+        # Build batch input in one preprocesser_call; avoid per-sample cat
+        batch_inputs = self.construct_model_input(
+            observations, prefix_list, postfix_list
+        )
+
+        padding_list = []
+        for ds_name in batch_inputs["dataset_names"]:
+            padding = (
+                torch.zeros_like(self.normalizer_action.delta[ds_name])
+                .unsqueeze(0)
+                .to("cpu")
+            )
+            padding_list.append(padding)
+        padding = torch.cat(padding_list, dim=0)
+        padding_action = self.normalizer_action.normalize_data(
+            padding, batch_inputs["dataset_names"]
+        ).to(batch_inputs["input_ids"].device)
+
+        with ScopeTimer("generate_flow_action_batch"):
+            model_output = self.model.generate_flow_action(
+                action_horizon=self.config.action_horizon,
+                action_dim=self.config.action_dim,
+                num_inference_timesteps=self.config.num_inference_timesteps,
+                padding_action=padding_action,
+                **batch_inputs,
+            )
+
+        predict_action = model_output["predict_action"]  # [B, H, D]
+
+        outputs = []
+        for i in range(batch_size):
+            single_action = predict_action[i : i + 1]
+            single_output = {
+                "predict_action": single_action,
+                "robot_state_action_data": observations[i]["robot_state_action_data"],
+            }
+            single_output["robot_state_action_data"].save_action_data(
+                single_output["predict_action"]
+            )
+            outputs.append(single_output)
+
+        return outputs
+
+    def infer_ar_action(self, observation, instruction):
+        self.logger.info("generating ar action")
+        self.logger.info(f"ar action instruction: {instruction}")
+
+        prefix_text, _ = self.get_text_for_action(instruction)
+        model_input = self.construct_model_input([observation], [prefix_text], [""])
+
+        model_output = self.model.generate_ar_action(
+            action_horizon=self.config.action_horizon,
+            action_dim=self.config.ar_action_dim,
+            num_inference_timesteps=self.config.num_inference_timesteps,
+            robot_type_id=self.robot_type_id,
+            **model_input,
+        )
+        self.logger.info("ar action generation done")
+
+        model_output["robot_state_action_data"] = observation["robot_state_action_data"]
+        model_output["robot_state_action_data"].save_action_data(
+            model_output["predict_action"]
+        )
+        self.logger.info("saved ar action to robot_state_action_data")
+        return model_output
+
+    def infer_subtask(self, observation, instruction):
+        self.logger.info("generating subtask")
+        self.logger.info(f"subtask instruction: {instruction}")
+
+        prefix_text, postfix_text = self.get_text_for_subtask_generation(instruction)
+        model_input = self.construct_model_input([observation], [prefix_text], [""])
+
+        model_output = self.model.generate_text(**model_input)
+        subtask = model_output["predict_output_text"][0].split("<|im_end|>")[0].strip()
+        self.logger.info(f"subtask generation done, subtask: {subtask}")
+        return subtask
+
+    def infer_vqa(self, observation, instruction):
+        self.logger.info("generating vqa answer")
+
+        if isinstance(observation["multi_modal"], list):
+            orig_size = observation["multi_modal"][0].size
+        else:
+            orig_size = observation["multi_modal"].size
+
+        prefix_text = instruction
+        model_input = self.construct_model_input([observation], [prefix_text], [""])
+
+        model_output = self.model.generate_text(**model_input)
+        answer = model_output["predict_output_text"][0].split("<|im_end|>")[0].strip()
+        self.logger.info(f"vqa answer done, answer: {answer}")
+        answer = reverse_grounding_points(
+            answer,
+            orig_size[1],
+            orig_size[0],
+            model_input["image_size"][0][1],
+            model_input["image_size"][0][0],
+            self.config.data_config.model_type,
+        )
+        points = extract_grounding_points(answer)
+
+        return {
+            "answer": answer,
+            "points": points,
+        }
+
+    def infer_dllm_action(
+        self, observation, instruction, use_ar_action=False, dataset_name="x2_normal"
+    ):
+        assert self.model_type in ["qwen2_5"], "DLLM only supports qwen2_5"
+        prefix_text, postfix_text = self.get_text_for_dllm_action(instruction)
+        model_input = self.construct_model_input(
+            [observation], [prefix_text], [postfix_text], pad_prefix=True
+        )
+        model_input = self.model.update_infer_dllm_position_mask(model_input)
+        total_ar_step = self.tokenizer_mixin.inference_ar_steps_for_dllm
+        cnt = 0
+        while cnt < 3:  # fast mode: at most 3 retries
+            model_output = self.model.generate_dllm_action(
+                action_horizon=self.config.action_horizon,
+                action_dim=self.config.action_dim,
+                ar_action_dim=self.config.ar_action_dim,
+                num_inference_timesteps=self.config.num_inference_timesteps,
+                use_ar_action=use_ar_action,
+                total_ar_step=total_ar_step,
+                robot_type_id=self.robot_type_id,  # for v3.1 delta decode
+                **model_input,
+            )
+            if model_output["predict_action"] is not None:
+                break
+            cnt += 1
+            self.logger.warning(f"dllm action generation failed, retry {cnt}")
+        self.logger.info("dllm action generation done")
+
+        model_output["robot_state_action_data"] = observation["robot_state_action_data"]
+        model_output["robot_state_action_data"].save_action_data(
+            model_output["predict_action"]
+        )
+        self.logger.info("saved dllm action to robot_state_action_data")
+        return model_output

wall_x/_vendor/harrix/serving/_wallx_infer/robot.py

@@ -312,9 +312,16 @@ class Robot(ABC):
         return state, views
 
     def _get_dof_mask(self):
-        dof_config = self.config.train_config["dof_config"]
+        from wall_x._vendor.harrix.utils.train_config import resolve_dof_config
+
+        dof_config = resolve_dof_config(self.config.train_config)
         total_dof = sum(dof_config.values())
         dof_mask = np.ones((1, self.config.action_horizon, total_dof))
+        start_idx = 0
+        for key, dof_size in dof_config.items():
+            if key == "action_padding":
+                dof_mask[:, :, start_idx : start_idx + dof_size] = 0
+            start_idx += dof_size
         return dof_mask
 
     @staticmethod

wall_x/_vendor/harrix/serving/launch_serving_rtc.py

@@ -0,0 +1,319 @@
+#!/usr/bin/env python3
+"""
+Server script for Wall-X model.
+
+This script serves a Wall-X model using a websocket server, allowing
+clients to connect and get action predictions from observations.
+
+"""
+
+import dataclasses
+import enum
+import inspect
+import logging
+import os
+import socket
+import sys
+import yaml
+import traceback
+import tyro
+
+from wall_x._vendor.harrix.serving.websocket_policy_server import WebsocketPolicyServer
+
+
+def _server_model_config_to_infer_kwargs(model_config) -> dict:
+    from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig
+
+    infer_params = inspect.signature(InferConfig.__init__).parameters
+    return {
+        k: v
+        for k, v in vars(model_config).items()
+        if v is not None and k in infer_params
+    }
+
+
+def get_wallx_policy(
+    model_config, image_passing_mode, serialize_actions=True, rtc_config=None
+):
+    from wall_x._vendor.harrix.serving.policy.wall_x_policy_rtc import WallXPolicy
+    from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig
+
+    config = InferConfig(**_server_model_config_to_infer_kwargs(model_config))
+    return WallXPolicy(
+        config=config,
+        image_passing_mode=image_passing_mode,
+        serialize_actions=serialize_actions,
+        rtc_config=rtc_config,
+    )
+
+
+
+
+logger = logging.getLogger(__name__)
+
+
+class EnvMode(enum.Enum):
+    """Supported environments/datasets."""
+
+    X2ROBOT = "x2robot"
+    LIBERO = "libero"
+
+
+@dataclasses.dataclass
+class ServerModelConfig:
+    """Configuration for loading a Wall-X model."""
+
+    checkpoint_path: str | None = None
+    train_config_path: str | None = None
+    # robot_host: str = '0.0.0.0'
+    # robot_port: int = 33723
+    robot_type: str = "desktop"  # ["desktop", "turtle"]
+    robot_action_start_ratio: float = (
+        0.0  # proportion of action execution to start from
+    )
+    robot_action_end_ratio: float = 1.0  # proportion of action execution to end at
+    robot_action_interpolate_multiplier: int = 10  # action interpolation multiplier
+    robot_use_joint_angle_control: bool = (
+        False  # use joint angle control (model must predict joints)
+    )
+    turtle_as_desktop: bool = (
+        False  # use turtle platform as desktop with fixed base/head/camera/height
+    )
+    action_horizon: int = 32  # specify the correct horizon for the model
+    action_dim: int | None = None
+    model_device: str = "cuda"
+    num_inference_timesteps: int = 10
+    num_inference_steps: int | None = None
+    cfg_scale: float | None = None
+    seed: int | None = None
+    save_video_dir: str = "./videos"
+
+    # Please specify explicitly if the checkpoint was not trained on the x2robot dataset.
+    norm_key: str | None = None
+    # Model cameras; None = infer from train config ``data.key_mappings.camera``.
+    cam_names: list[str] | None = None
+    # Robot camera keys in incoming websocket observations.
+    camera_front_key: str = "camera_front"
+    camera_left_key: str = "camera_left"
+    camera_right_key: str = "camera_right"
+    # Serving prompt controls. If the client request has no instruction,
+    # default_instruction is used. prompt_template follows train-config semantics.
+    default_instruction: str | None = None
+    prompt_template: str | None = None
+    qwen25_prompt_template: str | None = None
+    prompt_priority_order: str | None = None
+
+@dataclasses.dataclass
+class Args:
+    """Arguments for the serve_wall_x script."""
+
+    # Environment mode (used for default configurations)
+    env: EnvMode = EnvMode.X2ROBOT
+
+    # Model configuration. If not provided, uses default config for the environment
+    model_config: ServerModelConfig | None = None
+
+    # Default text prompt to use if not provided in observation
+    default_prompt: str | None = None
+
+    # Port to serve the policy on
+    port: int = 43007
+
+    # Host to bind the server to
+    host: str = "0.0.0.0"
+
+    # Enable debug logging
+    debug: bool = False
+
+    # Image passing mode
+    image_passing_mode: str = "base64"  # ["numpy", "base64"]
+
+    # Model type
+    model_type: str = "wallx"  # OSS supports qwen2.5 Wall-X only
+
+    # Serialize actions via robot_preprocessor (True for robot control, False for raw output)
+    serialize_actions: bool = True
+
+    # -- Dynamic batching -----------------------------------------
+    # Set max_batch_size to enable dynamic batching. None = single mode.
+    max_batch_size: int | None = None
+    max_wait_time_ms: float = 0
+    max_queue_size: int = 100
+    timeout_ms: float = 30000
+
+    # -- Real-Time Chunking ---------------------------------------
+    rtc_execution_horizon: int = 6
+    rtc_max_guidance_weight: float = 10.0
+    rtc_prefix_attention_schedule: str = "linear"
+
+    # -- Engine flags ---------------------------------------------
+    enable_experimental_engine: bool = False
+    enable_cuda_graph: bool = False
+
+
+# Default model configurations for each environment
+DEFAULT_CONFIGS: dict[EnvMode, ServerModelConfig] = {
+    EnvMode.X2ROBOT: ServerModelConfig(
+        checkpoint_path=None,
+        train_config_path=None,
+        robot_action_start_ratio=0.0,
+        robot_action_end_ratio=1.0,
+        robot_action_interpolate_multiplier=10,
+        robot_use_joint_angle_control=False,
+        turtle_as_desktop=False,
+        action_horizon=32,
+        action_dim=None,
+        model_device="cuda",
+        num_inference_timesteps=10,
+    ),
+    EnvMode.LIBERO: ServerModelConfig(
+        checkpoint_path=None,
+        train_config_path=None,
+        robot_type="desktop",
+        robot_action_start_ratio=0.0,
+        robot_action_end_ratio=1.0,
+        robot_action_interpolate_multiplier=1,
+        action_horizon=10,
+        action_dim=None,
+        model_device="cuda",
+        num_inference_timesteps=10,
+        cam_names=None,
+    ),
+}
+
+
+def get_model_config(args: Args) -> ServerModelConfig:
+    """Get model configuration from args or defaults."""
+    if args.model_config is not None:
+        return args.model_config
+
+    if config := DEFAULT_CONFIGS.get(args.env):
+        logger.info(f"Using default configuration for {args.env.value}")
+        return config
+
+    raise ValueError(
+        f"No default configuration for {args.env.value}. "
+        f"Please provide --model-config with model_path and action_tokenizer_path."
+    )
+
+
+def create_policy(args: Args):
+    """Create a policy from the given arguments."""
+    model_config = get_model_config(args)
+    if args.model_type != "wallx":
+        raise ValueError(
+            f"Unsupported model type: {args.model_type!r}. "
+            "The public package only supports model_type='wallx'."
+        )
+    from wall_x._vendor.harrix.serving.rtc_wallx import WallXRTCConfig
+
+    if args.max_batch_size is not None:
+        raise ValueError(
+            "RTC serving keeps per-session chunk state and does not support "
+            "dynamic batching; omit --max-batch-size."
+        )
+    rtc_config = WallXRTCConfig(
+        execution_horizon=args.rtc_execution_horizon,
+        max_guidance_weight=args.rtc_max_guidance_weight,
+        prefix_attention_schedule=args.rtc_prefix_attention_schedule,
+    )
+    policy = get_wallx_policy(
+        model_config,
+        args.image_passing_mode,
+        args.serialize_actions,
+        rtc_config=rtc_config,
+    )
+    return policy
+
+
+def main(args: Args) -> None:
+    """Main function to start the Wall-X model server."""
+    log_level = logging.DEBUG if args.debug else logging.INFO
+    logging.basicConfig(
+        level=log_level,
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    )
+
+    # Set engine environment variables
+    if args.enable_experimental_engine:
+        os.environ["ENABLE_EXPERIMENTAL_INFERENCE_ENGINE"] = "true"
+        logger.info("ENABLE_EXPERIMENTAL_INFERENCE_ENGINE=true")
+
+    if args.enable_cuda_graph:
+        os.environ["ENABLE_CUDA_GRAPH"] = "true"
+        logger.info("ENABLE_CUDA_GRAPH=true")
+
+    logger.info("Starting model server")
+    logger.info(f"Model type: {args.model_type}")
+    logger.info(f"Environment: {args.env.value}")
+    logger.info(f"Port: {args.port}")
+    logger.info(f"Host: {args.host}")
+    logger.info(f"Serialize actions: {args.serialize_actions}")
+    logger.info(
+        "RTC: execution_horizon=%d max_guidance_weight=%.3f schedule=%s",
+        args.rtc_execution_horizon,
+        args.rtc_max_guidance_weight,
+        args.rtc_prefix_attention_schedule,
+    )
+
+    # Create policy
+    try:
+        policy = create_policy(args)
+    except Exception as e:
+        logger.error(f"Failed to create policy: {e}")
+        logger.error(traceback.format_exc())
+        sys.exit(1)
+
+    # Get policy metadata
+    policy_metadata = policy.metadata
+    policy_metadata["env"] = args.env.value
+    policy_metadata["rtc"] = {
+        "enabled": True,
+        "execution_horizon": args.rtc_execution_horizon,
+        "max_guidance_weight": args.rtc_max_guidance_weight,
+        "prefix_attention_schedule": args.rtc_prefix_attention_schedule,
+    }
+
+    # Get network info
+    hostname = socket.gethostname()
+    try:
+        local_ip = socket.gethostbyname(hostname)
+    except Exception:
+        local_ip = "unknown"
+
+    logger.info(f"Server hostname: {hostname}")
+    logger.info(f"Server IP: {local_ip}")
+    logger.info(f"Server will be available at: ws://{args.host}:{args.port}")
+    logger.info(f"Health check endpoint: http://{args.host}:{args.port}/healthz")
+
+    batching_str = (
+        f"batch_size={args.max_batch_size}, wait={args.max_wait_time_ms}ms"
+        if args.max_batch_size
+        else "disabled"
+    )
+    logger.info(f"Batching: {batching_str}")
+
+    # Create and start server
+    server = WebsocketPolicyServer(
+        policy=policy,
+        host=args.host,
+        port=args.port,
+        metadata=policy_metadata,
+        max_batch_size=args.max_batch_size,
+        max_wait_time_ms=args.max_wait_time_ms,
+        max_queue_size=args.max_queue_size,
+        timeout_ms=args.timeout_ms,
+    )
+
+    logger.info("Starting server...")
+    try:
+        server.serve_forever()
+    except KeyboardInterrupt:
+        logger.info("Server stopped by user")
+    except Exception as e:
+        logger.error(f"Server error: {e}")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main(tyro.cli(Args))

wall_x/_vendor/harrix/serving/policy/wall_x_policy.py

@@ -128,6 +128,18 @@ class WallXPolicy(BasePolicy):
         response: Dict[str, Any] = {
             "predict_action": np.asarray(predict_action, dtype=np.float32),
         }
+        # Single-arm checkpoints normally return raw model output because the
+        # generic serializer expects both arms.  Still expose the official
+        # right-arm reconstruction (absolute xyz/euler/gripper) so legacy
+        # clients do not have to guess relative-action semantics.
+        state_action = model_output.get("robot_state_action_data")
+        if state_action is not None and self._is_single_arm_right_only():
+            try:
+                response["follow2_pos"] = self.robot_preprocessor._get_right_arm_action(
+                    state_action
+                ).astype(np.float32)
+            except Exception as exc:
+                logger.warning("Could not serialize single-arm right action: %s", exc)
         if "subtask" in model_output:
             response["subtask"] = model_output["subtask"]
         return response

wall_x/_vendor/harrix/serving/policy/wall_x_policy_rtc.py

@@ -0,0 +1,472 @@
+import base64
+import logging
+import threading
+from typing import Dict, Any, List
+
+import cv2
+import numpy as np
+import torch
+from wall_x._vendor.harrix.serving.websocket_policy_server import BasePolicy
+from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig
+from wall_x._vendor.harrix.serving._wallx_infer.model_wrapper_rtc import WallxModelWrapper
+from wall_x._vendor.harrix.serving._wallx_infer.robot import (
+    DesktopRobotPreprocessor,
+    TurtleRobotPreprocessor,
+    EX001RobotPreprocessor,
+)
+from wall_x._vendor.harrix.serving.policy._smoothing import apply_smoothing
+from wall_x.utils.timers import ScopeTimer
+from wall_x._vendor.harrix.serving.rtc_wallx import WallXRTCConfig
+from wall_x._vendor.x2robot_utils import geometry as geom
+
+logger = logging.getLogger(__name__)
+
+# Inference modes aligned with model_wrapper for vqa / batch extensions
+INFER_MODE_FLOW = "flow"
+INFER_MODE_AR = "ar"
+INFER_MODE_FLOW_WITH_SUBTASK = "flow_with_subtask"
+INFER_MODE_DLLM_FLOW = "dllm"
+INFER_MODE_DLLM_DD = "discrete_diffusion"
+ACTION_INFER_MODES = (
+    INFER_MODE_FLOW,
+    INFER_MODE_AR,
+    INFER_MODE_FLOW_WITH_SUBTASK,
+    INFER_MODE_DLLM_FLOW,
+    INFER_MODE_DLLM_DD,
+)
+
+
+class WallXPolicy(BasePolicy):
+    """Policy wrapper for Wall-X model that implements the BasePolicy interface."""
+
+    def __init__(
+        self,
+        config: InferConfig,
+        image_passing_mode: str = "base64",
+        default_infer_mode: str = INFER_MODE_FLOW,
+        serialize_actions: bool = True,
+        rtc_config: WallXRTCConfig | None = None,
+    ):
+        """Initialize the Wall-X policy.
+
+        Args:
+            config: Inference configuration dataclass.
+            image_passing_mode: How images are passed from client ('base64' or 'numpy').
+            default_infer_mode: Default inference mode ('flow', 'ar', 'flow_with_subtask', etc.).
+            serialize_actions: If True, actions are serialized via robot_preprocessor;
+                if False, raw model output is returned directly.
+        """
+        self.config = config
+        self.rtc_config = rtc_config or WallXRTCConfig()
+        self.model_wrapper = WallxModelWrapper(config, rtc_config=self.rtc_config)
+        self.robot_preprocessor = self._register_robot_preprocessor()
+        self._rtc_sessions: dict[str, dict[str, Any]] = {}
+        self._rtc_infer_lock = threading.Lock()
+        self.image_passing_mode = image_passing_mode
+        self.default_infer_mode = default_infer_mode
+        self.serialize_actions = serialize_actions
+        logger.info(
+            "Image passing mode: %s, robot_type: %s, default_infer_mode: %s, "
+            "serialize_actions: %s, smooth_action: %s, smooth_gripper: %s",
+            self.image_passing_mode,
+            config.robot_type,
+            self.default_infer_mode,
+            self.serialize_actions,
+            config.smooth_action,
+            config.smooth_gripper,
+        )
+
+    def _register_robot_preprocessor(self):
+        """Select preprocessor by config.robot_type; mirrors env._register_robot."""
+        if self.config.robot_type == "desktop":
+            return DesktopRobotPreprocessor(self.config)
+        if self.config.robot_type == "turtle":
+            return TurtleRobotPreprocessor(self.config)
+        if self.config.robot_type == "ex001":
+            return EX001RobotPreprocessor(self.config)
+        raise ValueError(f"Invalid robot_type: {self.config.robot_type!r}")
+
+    @staticmethod
+    def _decode_base64_image(image_b64: str) -> np.ndarray:
+        """Decode client base64 JPEG/PNG payloads to RGB images."""
+        img_bytes = base64.b64decode(image_b64)
+        img_array = np.frombuffer(img_bytes, dtype=np.uint8)
+        decoded_img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
+        if decoded_img is None:
+            raise ValueError("Failed to decode base64 image payload")
+        return cv2.cvtColor(decoded_img, cv2.COLOR_BGR2RGB)
+
+    def reset(self) -> None:
+        """Reset the policy state."""
+        self.action_buffer = []
+        self.buffer_index = 0
+        logger.debug("Policy reset")
+
+    def _get_dof_config(self) -> dict:
+        train_config = self.config.train_config or {}
+        return (
+            train_config.get("dof_config")
+            or train_config.get("task", {}).get("dof_config", {})
+            or {}
+        )
+
+    def _is_single_arm_right_only(self) -> bool:
+        """True for single-arm LIBERO-style configs (right arm only, no left arm)."""
+        train_config = self.config.train_config or {}
+        agent_pos = (
+            train_config.get("agent_pos_config")
+            or train_config.get("task", {}).get("agent_pos_config")
+            or {}
+        )
+        skip = {"action_padding"}
+        has_left = any(k.startswith("follow_left_") for k in agent_pos if k not in skip)
+        has_right = any(
+            k.startswith("follow_right_") for k in agent_pos if k not in skip
+        )
+        return has_right and not has_left
+
+    def _pack_action_chunk_response(
+        self, model_output: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Return a msgpack-safe action chunk for websocket clients."""
+        predict_action = model_output["predict_action"]
+        if isinstance(predict_action, torch.Tensor):
+            predict_action = predict_action.detach().cpu().numpy()
+        response: Dict[str, Any] = {
+            "predict_action": np.asarray(predict_action, dtype=np.float32),
+        }
+        # Single-arm checkpoints normally return raw model output because the
+        # generic serializer expects both arms.  Still expose the official
+        # right-arm reconstruction (absolute xyz/euler/gripper) so legacy
+        # clients do not have to guess relative-action semantics.
+        state_action = model_output.get("robot_state_action_data")
+        if state_action is not None and self._is_single_arm_right_only():
+            try:
+                response["follow2_pos"] = self.robot_preprocessor._get_right_arm_action(
+                    state_action
+                ).astype(np.float32)
+            except Exception as exc:
+                logger.warning("Could not serialize single-arm right action: %s", exc)
+        if "subtask" in model_output:
+            response["subtask"] = model_output["subtask"]
+        return response
+
+    def _get_predict_action_keys(self) -> List[str]:
+        """Resolve flat action key order for smoothing / serialization."""
+        data_cfg = self.config.data_config
+        keys = None
+        if isinstance(data_cfg, dict):
+            keys = data_cfg.get("predict_action_keys")
+        else:
+            keys = getattr(data_cfg, "predict_action_keys", None)
+        if keys:
+            return list(keys)
+        return list(self._get_dof_config().keys())
+
+    def _apply_smoothing(self, model_output: Dict[str, Any]) -> None:
+        """Apply Laplacian smoothing on predict_action before 6D->euler conversion."""
+        cfg = self.config
+        if not getattr(cfg, "smooth_action", False):
+            return
+        dof_config = self._get_dof_config()
+        apply_smoothing(
+            model_output,
+            smooth_action=True,
+            smooth_gripper=cfg.smooth_gripper,
+            predict_action_keys=self._get_predict_action_keys(),
+            action_padding_dof=dof_config.get("action_padding"),
+            action_dim=cfg.action_dim,
+        )
+
+    def _rtc_build_normalized_prefix(self, state: Dict, rtc: Dict):
+        session_id = str(rtc.get("session_id", "default"))
+        if rtc.get("reset"):
+            self._rtc_sessions.pop(session_id, None)
+            return None
+        cached = self._rtc_sessions.get(session_id)
+        if not cached:
+            return None
+
+        consumed = max(0, int(rtc.get("consumed_model_steps", 0)))
+        absolute = np.asarray(cached["follow2_pos"], dtype=np.float64)
+        if consumed >= len(absolute):
+            return None
+        absolute = absolute[consumed:]
+        current = np.asarray(state["follow2_pos"], dtype=np.float64).reshape(7)
+
+        state_rotation = geom.euler_to_matrix_zyx_batch_nb(
+            current[None, 3:6]
+        )[0]
+        absolute_rotation = geom.euler_to_matrix_zyx_batch_nb(absolute[:, 3:6])
+        delta_rotation = absolute_rotation @ state_rotation.T
+        delta_rotation_6d = delta_rotation[:, :2, :].reshape(len(absolute), 6)
+
+        dof_config = self._get_dof_config()
+        columns = []
+        for key, width in dof_config.items():
+            width = int(width)
+            lowered = key.lower()
+            if key == "action_padding":
+                values = np.zeros((len(absolute), width), dtype=np.float64)
+            elif "follow_right" in lowered and "cartesian_pos" in lowered:
+                values = absolute[:, :3] - current[None, :3]
+            elif "follow_right" in lowered and "rotation_6d" in lowered:
+                values = delta_rotation_6d
+            elif "follow_right" in lowered and "rotation" in lowered:
+                raise ValueError("RTC currently requires the checkpoint's 6D rotation layout")
+            elif "follow_right" in lowered and "gripper" in lowered:
+                values = absolute[:, 6:7]
+            else:
+                values = np.zeros((len(absolute), width), dtype=np.float64)
+            if values.shape[1] != width:
+                raise ValueError(
+                    f"RTC field {key!r} expected width {width}, got {values.shape[1]}"
+                )
+            columns.append(values)
+
+        raw = np.concatenate(columns, axis=1)
+        tensor = torch.from_numpy(raw).to(
+            device=self.config.model_device, dtype=torch.float32
+        ).unsqueeze(0)
+        return self.model_wrapper.normalizer_action.normalize_data(
+            tensor, [self.model_wrapper.norm_key]
+        )
+
+    def _rtc_update_session(
+        self,
+        *,
+        state: Dict,
+        rtc: Dict,
+        model_output: Dict[str, Any],
+    ) -> None:
+        session_id = str(rtc.get("session_id", "default"))
+        state_action = model_output.get("robot_state_action_data")
+        if state_action is None:
+            return
+        absolute = self.robot_preprocessor._get_right_arm_action(state_action)
+        absolute = np.asarray(absolute, dtype=np.float64)
+        if len(absolute) == self.config.action_horizon + 1:
+            absolute = absolute[1:]
+        self._rtc_sessions[session_id] = {
+            "follow2_pos": absolute.copy(),
+            "chunk_id": int(rtc.get("request_id", 0)),
+        }
+
+    def _run_action_infer(
+        self,
+        observation: Dict,
+        instruction: str,
+        mode: str,
+        *,
+        rtc_context: Dict | None = None,
+        rtc_prefix=None,
+    ) -> Dict[str, Any]:
+        """Run model_wrapper action inference by mode; returns model_output with robot_state_action_data.
+
+        Supports flow | ar | flow_with_subtask. VQA can be added here later (different return format).
+        """
+        if mode == INFER_MODE_FLOW:
+            rtc_context = rtc_context or {}
+            return self.model_wrapper.infer_flow_action(
+                observation,
+                instruction,
+                prev_chunk_left_over=rtc_prefix,
+                inference_delay=int(rtc_context.get("inference_delay_steps", 0)),
+                execution_horizon=int(
+                    rtc_context.get(
+                        "execution_horizon",
+                        self.rtc_config.execution_horizon,
+                    )
+                ),
+            )
+        if mode == INFER_MODE_AR:
+            return self.model_wrapper.infer_ar_action(observation, instruction)
+        if mode == INFER_MODE_FLOW_WITH_SUBTASK:
+            with ScopeTimer("infer_subtask"):
+                subtask = self.model_wrapper.infer_subtask(observation, instruction)
+            with ScopeTimer("infer_flow_action"):
+                model_output = self.model_wrapper.infer_flow_action(
+                    observation, subtask
+                )
+            model_output["subtask"] = subtask
+            return model_output
+        if mode == INFER_MODE_DLLM_FLOW:
+            return self.model_wrapper.infer_dllm_action(
+                observation, instruction, use_ar_action=False
+            )
+        if mode == INFER_MODE_DLLM_DD:
+            return self.model_wrapper.infer_dllm_action(
+                observation, instruction, use_ar_action=True
+            )
+        raise ValueError(
+            f"Unsupported infer_mode={mode!r}, expected one of {ACTION_INFER_MODES}"
+        )
+
+    def infer(self, obs: Dict) -> Dict:
+        with self._rtc_infer_lock:
+            return self._infer_locked(obs)
+
+    def _infer_locked(self, obs: Dict) -> Dict:
+        """Infer action from observation.
+
+        Args:
+            obs: Dictionary containing:
+                - 'state': Robot state
+                - 'views': Camera views (keyed by camera name)
+                - 'instruction': Task instruction
+                - Optional: 'infer_mode' - one of 'flow' | 'ar' | 'flow_with_subtask'
+                - Optional: 'robot_action_start_ratio' / 'robot_action_end_ratio' /
+                  'robot_action_interpolate_multiplier' - override config for action trim/interpolate
+
+        Returns:
+            When serialize_actions=True (default): Serialized action dict
+                (e.g. follow1_pos/follow2_pos or follow1_joints/follow2_joints).
+            When serialize_actions=False: Raw model_output dict.
+            If infer_mode is flow_with_subtask, also includes 'subtask'.
+        """
+        state = obs["state"]
+        views = obs["views"]
+        instruction = obs.get("instruction") or self.config.default_instruction or ""
+        if self.image_passing_mode == "base64":
+            for k, v in views.items():
+                views[k] = np.expand_dims(self._decode_base64_image(v), axis=0)
+
+        with ScopeTimer("get_observation"):
+            observation = self.robot_preprocessor.get_observation(state, views)
+
+        mode = obs.get("infer_mode", self.default_infer_mode)
+        rtc_context = dict(obs.get("rtc") or {})
+        rtc_prefix = self._rtc_build_normalized_prefix(state, rtc_context)
+        with ScopeTimer(f"infer_{mode}"):
+            model_output = self._run_action_infer(
+                observation,
+                instruction,
+                mode,
+                rtc_context=rtc_context,
+                rtc_prefix=rtc_prefix,
+            )
+
+        self._apply_smoothing(model_output)
+        if (
+            "robot_state_action_data" in model_output
+            and "predict_action" in model_output
+        ):
+            model_output["robot_state_action_data"].save_action_data(
+                model_output["predict_action"]
+            )
+
+        self._rtc_update_session(
+            state=state,
+            rtc=rtc_context,
+            model_output=model_output,
+        )
+
+        # Single-arm LIBERO has no left-arm EE to serialize; dual-arm follow1/follow2
+        # layout would fail in get_serialized_actions. Return raw action chunks instead.
+        if not self.serialize_actions or self._is_single_arm_right_only():
+            response = self._pack_action_chunk_response(model_output)
+        else:
+            response = self.robot_preprocessor.get_serialized_actions(
+                model_output, robot_action_interpolate_multiplier=1
+            )
+        response["_rtc"] = {
+            "session_id": str(rtc_context.get("session_id", "default")),
+            "request_id": int(rtc_context.get("request_id", 0)),
+            "guided": rtc_prefix is not None,
+            "inference_delay_steps": int(
+                rtc_context.get("inference_delay_steps", 0)
+            ),
+            "execution_horizon": int(
+                rtc_context.get(
+                    "execution_horizon",
+                    self.rtc_config.execution_horizon,
+                )
+            ),
+        }
+        return response
+
+    # -- Batch inference ----------------------------------------------
+
+    def _preprocess_obs(self, obs: Dict):
+        """Preprocess a single observation dict into (observation, instruction).
+
+        Handles both base64 and raw image modes.
+        """
+        state = obs["state"]
+        views = obs["views"]
+        instruction = obs.get("instruction") or self.config.default_instruction or ""
+
+        if self.image_passing_mode == "base64":
+            for k, v in views.items():
+                views[k] = np.expand_dims(self._decode_base64_image(v), axis=0)
+        else:
+            # raw numpy: ensure (1, H, W, C)
+            for k, img in views.items():
+                if isinstance(img, np.ndarray) and img.ndim == 3:
+                    views[k] = np.expand_dims(img, axis=0)
+
+        observation = self.robot_preprocessor.get_observation(state, views)
+        return observation, instruction
+
+    def infer_batch(
+        self, obs_list: List[Dict[str, Any]], skip_serialize: bool = False
+    ) -> List[Dict[str, Any]]:
+        """Perform batch inference.
+
+        Args:
+            obs_list: List of observations, each containing:
+                - "views": dict of camera images
+                - "state": robot state dict or array
+                - "instruction": text instruction
+
+        Returns:
+            List of action dicts.
+        """
+        batch_size = len(obs_list)
+        logger.info(f"WallXPolicy.infer_batch: processing {batch_size} observations")
+
+        try:
+            observations = []
+            instructions = []
+            for obs in obs_list:
+                observation, instruction = self._preprocess_obs(obs)
+                observations.append(observation)
+                instructions.append(instruction)
+
+            mode = obs_list[0].get("infer_mode", self.default_infer_mode)
+
+            with torch.no_grad():
+                if mode == INFER_MODE_FLOW:
+                    model_outputs = self.model_wrapper.infer_flow_action_batch(
+                        observations, instructions
+                    )
+                else:
+                    # AR / flow_with_subtask: fall back to per-sample inference
+                    model_outputs = []
+                    for obs_dict, instruction in zip(observations, instructions):
+                        output = self._run_action_infer(obs_dict, instruction, mode)
+                        model_outputs.append(output)
+
+            if skip_serialize:
+                return [{}] * len(model_outputs)
+
+            results = []
+            for model_output in model_outputs:
+                self._apply_smoothing(model_output)
+                if self.serialize_actions and not self._is_single_arm_right_only():
+                    action = self.robot_preprocessor.get_serialized_actions(
+                        model_output, robot_action_interpolate_multiplier=1
+                    )
+                    results.append(action)
+                else:
+                    results.append(self._pack_action_chunk_response(model_output))
+
+            return results
+
+        except Exception as e:
+            logger.error(f"Batch inference failed: {e}", exc_info=True)
+            return [{"action": {}, "error": str(e)} for _ in obs_list]
+
+    @property
+    def metadata(self) -> Dict[str, Any]:
+        return {"batch_enabled": True, "model": "wall-x"}

wall_x/_vendor/harrix/serving/rtc_wallx.py

@@ -0,0 +1,119 @@
+"""Wall-X Real-Time Chunking helpers.
+
+The guided update follows LeRobot's RTC implementation, adapted to Wall-X's
+flow convention: Wall-X integrates normalized actions from noise at t=0 to a
+clean action at t=1 with velocity ``dx/dt``.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import math
+
+import torch
+
+
+@dataclass(frozen=True)
+class WallXRTCConfig:
+    enabled: bool = True
+    execution_horizon: int = 6
+    max_guidance_weight: float = 10.0
+    prefix_attention_schedule: str = "linear"
+
+    def __post_init__(self):
+        if self.execution_horizon <= 0:
+            raise ValueError("execution_horizon must be positive")
+        if self.max_guidance_weight <= 0:
+            raise ValueError("max_guidance_weight must be positive")
+        if self.prefix_attention_schedule not in {"zeros", "ones", "linear", "exp"}:
+            raise ValueError(
+                "prefix_attention_schedule must be zeros, ones, linear, or exp"
+            )
+
+
+class WallXRTCProcessor:
+    """Inference-time RTC guidance in normalized Wall-X action space."""
+
+    def __init__(self, config: WallXRTCConfig):
+        self.config = config
+
+    def get_prefix_weights(self, start: int, end: int, total: int) -> torch.Tensor:
+        start = max(0, min(int(start), int(end), int(total)))
+        end = max(start, min(int(end), int(total)))
+        schedule = self.config.prefix_attention_schedule
+        if schedule == "zeros":
+            weights = torch.zeros(total)
+            weights[:start] = 1.0
+            return weights
+        if schedule == "ones":
+            weights = torch.zeros(total)
+            weights[:end] = 1.0
+            return weights
+
+        middle_len = end - start
+        if middle_len:
+            middle = torch.linspace(1.0, 0.0, middle_len + 2)[1:-1]
+            if schedule == "exp":
+                middle = middle * torch.expm1(middle) / (math.e - 1.0)
+        else:
+            middle = torch.empty(0)
+        return torch.cat(
+            [torch.ones(start), middle, torch.zeros(total - end)], dim=0
+        )
+
+    def guide_increasing_flow(
+        self,
+        *,
+        x_t: torch.Tensor,
+        time: torch.Tensor | float,
+        predict_velocity,
+        prev_chunk_left_over: torch.Tensor | None,
+        inference_delay: int = 0,
+        execution_horizon: int | None = None,
+    ) -> torch.Tensor:
+        """Return RTC-guided velocity for a flow integrated from t=0 to t=1."""
+        if prev_chunk_left_over is None or not self.config.enabled:
+            return predict_velocity(x_t)
+
+        x = x_t.detach().clone().requires_grad_(True)
+        prefix = prev_chunk_left_over.to(device=x.device, dtype=x.dtype)
+        if prefix.ndim == 2:
+            prefix = prefix.unsqueeze(0)
+        if prefix.shape[0] == 1 and x.shape[0] > 1:
+            prefix = prefix.expand(x.shape[0], -1, -1)
+
+        padded = torch.zeros_like(x)
+        steps = min(prefix.shape[1], x.shape[1])
+        dims = min(prefix.shape[2], x.shape[2])
+        padded[:, :steps, :dims] = prefix[:, :steps, :dims]
+
+        horizon = execution_horizon or self.config.execution_horizon
+        horizon = min(int(horizon), steps, x.shape[1])
+        weights = self.get_prefix_weights(inference_delay, horizon, x.shape[1])
+        weights = weights.to(device=x.device, dtype=x.dtype).view(1, -1, 1)
+
+        with torch.enable_grad():
+            velocity = predict_velocity(x)
+            t = torch.as_tensor(time, device=x.device, dtype=x.dtype)
+            remaining = torch.clamp(1.0 - t, min=1e-6)
+            clean_estimate = x + remaining * velocity
+            error = (padded - clean_estimate) * weights
+            correction = torch.autograd.grad(
+                clean_estimate,
+                x,
+                grad_outputs=error.detach(),
+                retain_graph=False,
+            )[0]
+
+        # Same guidance schedule as LeRobot RTC after mapping its 1->0 time
+        # convention to Wall-X's 0->1 convention.
+        t_safe = torch.clamp(t, min=1e-6)
+        inv_r2 = (remaining.square() + t.square()) / remaining.square()
+        weight = (remaining / t_safe) * inv_r2
+        weight = torch.nan_to_num(
+            weight,
+            nan=self.config.max_guidance_weight,
+            posinf=self.config.max_guidance_weight,
+        ).clamp(max=self.config.max_guidance_weight)
+        return (velocity + weight * correction).detach()
+