import yaml import os from wall_x.model.model_utils import update_model_config # from x2robot_dataset.configs.config import X2RDataConfig import json from typing import List, Dict, Optional, Any from dataclasses import dataclass, field @dataclass class X2RDataConfig: """ Unified X2Robot data configuration class (reorganized by README's 5 modules): 1) Data I/O and caching 2) Visual input and sampling (image/camera) 3) Action and time series 4) Instruction and multimodal 5) Data cleaning and alignment (validation/augmentation/framework constraints) """ # ---------------------------------------------------------------------- # 1) Data I/O and caching # ---------------------------------------------------------------------- cache_dir: str = "~/.cache/dataset_cache" dataset_config_path: Optional[str] = None use_cache: bool = True check_mode: bool = True preload_size: int = 128 buffer_size: int = 20000 batch_size: int = 32 train_test_split: float = 0.9 seed: int = 42 episode_chunk_size: int = ( 500 # Commonly used on VG side (number of frames for episode chunking) ) # ---------------------------------------------------------------------- # 2) Visual input and sampling (image/camera) # ---------------------------------------------------------------------- # Camera mapping cam_mapping: Dict[str, str] = field( default_factory=lambda: { "faceImg": "face_view", "leftImg": "left_wrist_view", "rightImg": "right_wrist_view", } ) # Image and augmentation resolution: Dict[str, int] = field( default_factory=lambda: { "face_view": -1, "left_wrist_view": 128, "right_wrist_view": 128, } ) cam_augmentation_list: List[str] = field(default_factory=list) # Image time series (history/future) image_horizon: int = 1 image_history_length: int = 0 image_history_interval: int = 1 future_image_length: int = 0 future_image_interval: int = 1 future_image_indices: Optional[List[int]] = ( None # If provided, length must equal image_horizon ) # Smart scaling max_pixels: int = field( default_factory=lambda: 1280 * 28 * 28 ) # Will be replaced with MAX_PIXELS in __post_init__ min_pixels: int = field( default_factory=lambda: 4 * 28 * 28 ) # Will be replaced with MIN_PIXELS in __post_init__ image_factor: int = 28 # Will be replaced with IMAGE_FACTOR in __post_init__ # ---------------------------------------------------------------------- # 3) Action and time series # ---------------------------------------------------------------------- predict_action_keys: List[str] = field(default_factory=list) obs_action_keys: List[str] = field(default_factory=list) # Action window action_horizon: int = 21 action_history_length: int = 0 action_horizon_flow: int = 32 action_horizon_ar: int = 0 # Padding strategy left_padding: bool = True right_padding: bool = True # Dimension configuration dof_config: Dict[str, int] = field(default_factory=dict) # Input degrees of freedom agent_pos_config: Dict[str, int] = field( default_factory=dict ) # Output degrees of freedom # State augmentation state_augmentation_ratio: float = 1.0 # Ratio of augmented states state_augmentation_prob: float = ( 0.1 # Random dimension masking probability for state string ) state_drop_prob: float = 0.0 # Probability of dropping entire state # ---------------------------------------------------------------------- # 4) Instruction and multimodal # ---------------------------------------------------------------------- default_instruction: str = "" instruction_path: Optional[str] = None instruction_key: Optional[List[Dict]] = None multimodal_chunk_size: int = 500 generate_subtask_ratio: float = 0.0 cot_ratio: float = 0.0 multimodal_data_ratio: float = ( 0.25 # Multimodal data ratio per batch in VLA dataset ) instruction_key_prob: Optional[Dict[str, float]] = None trunc_action_with_instruction: bool = True use_embodied_system_prompt_ratio: float = 0.0 # ---------------------------------------------------------------------- # 5) Data cleaning and alignment (validation/augmentation/framework constraints) # ---------------------------------------------------------------------- filter_angle_outliers: bool = False trim_stationary: bool = False use_state_string_representation: bool = False pad_prefix_to_same_length: bool = False put_ar_predict_in_postfix: bool = ( False # Whether to put ar prediction in postfix, set to True in prediction mode, False in training ) pad_to_128_multiple: bool = ( False # Triton Attention requirement (deprecated, always set to False) ) max_seqlen: int = 768 model_type: Optional[str] = None # qwen2_5, qwen2 model_config_path: Optional[str] = ( None # Model config path (used to derive PaddingSide) ) low_dim_obs_horizon: int = 1 # To be deprecated # ---------------------------------------------------------------------- # Validation and post-processing # ---------------------------------------------------------------------- def __post_init__(self): # TODO: Determine VGA model type validation here # assert self.model_type in ["qwen2_5", "qwen3"], f"Unsupported model type: {self.model_type}" if self.model_type == "qwen2_5": self.max_pixels = 16384 * 28 * 28 self.min_pixels = 4 * 28 * 28 self.image_factor = 28 elif self.model_type == "qwen3": self.max_pixels = 16384 * 32 * 32 self.min_pixels = 4 * 32 * 32 self.image_factor = 32 # Future image indices validation if ( self.future_image_indices and len(self.future_image_indices) != self.image_horizon ): raise ValueError( f"future_image_indices length must equal image_horizon: " f"{len(self.future_image_indices)} != {self.image_horizon}" ) # Auto-derive action window if self.action_horizon == 0: self.action_horizon = max(self.action_horizon_flow, self.action_horizon_ar) # Auto-derive action keys if not self.obs_action_keys: self.obs_action_keys = list(self.agent_pos_config.keys()) if not self.predict_action_keys: self.predict_action_keys = list(self.dof_config.keys()) # Derive PaddingSide # @Ryan: Only FlashAttention can use RightPadding, other AttnImpl use LeftPadding if self.model_config_path is not None: with open(self.model_config_path, "r", encoding="utf-8") as f: cfg = json.load(f) attn_impl = cfg["_attn_implementation"] if attn_impl == "flash_attention_2": self.padding_side = "right" else: self.padding_side = "left" # Convenience properties @property def use_6D_rotation(self) -> bool: """Whether to use 6D rotation (auto-determined from predict_action_keys)""" if hasattr(self, "_use_6D_rotation"): return self._use_6D_rotation self._use_6D_rotation = any("6D" in key for key in self.predict_action_keys) return self._use_6D_rotation @property def use_relative_action(self) -> bool: """Whether to use relative action (auto-determined from predict_action_keys)""" if hasattr(self, "_use_relative_action"): return self._use_relative_action self._use_relative_action = any( "relative" in key for key in self.predict_action_keys ) return self._use_relative_action # ---------------------------------------------------------------------- # YAML initialization # ---------------------------------------------------------------------- @classmethod def from_yaml_dict(cls, yaml_dict: Dict[str, Any]) -> "X2RDataConfig": """ Create config object from YAML config dict. Prioritizes data sub-config, then top-level fields. """ data_config = yaml_dict.get("data", {}) params: Dict[str, Any] = {} # 1) Data I/O and caching params.update( { "cache_dir": data_config.get( "cache_dir", yaml_dict.get("cache_dir", "~/.cache/dataset_cache") ), "dataset_config_path": data_config.get( "dataset_config_path", yaml_dict.get("dataset_config_path", None) ), "use_cache": data_config.get( "use_cache", yaml_dict.get("use_cache", True) ), "check_mode": data_config.get( "check_mode", yaml_dict.get("check_mode", True) ), "preload_size": data_config.get( "preload_size", yaml_dict.get("preload_size", 128) ), "buffer_size": data_config.get( "buffer_size", yaml_dict.get("buffer_size", 20000) ), "batch_size": data_config.get( "batch_size", yaml_dict.get( "batch_size_per_gpu", yaml_dict.get("batch_size", 32) ), ), "train_test_split": data_config.get("train_test_split", 0.9), "seed": yaml_dict.get("seed", 42), "episode_chunk_size": data_config.get("episode_chunk_size", 500), } ) # 2) Visual input and sampling (image/camera) params.update( { "cam_mapping": data_config.get( "cam_mapping", { "faceImg": "face_view", "leftImg": "left_wrist_view", "rightImg": "right_wrist_view", }, ), "resolution": data_config.get( "resolution", {"face_view": -1, "left_wrist_view": 128, "right_wrist_view": 128}, ), "cam_augmentation_list": data_config.get("cam_augmentation_list", []), "image_horizon": data_config.get("image_horizon", 1), "image_history_length": data_config.get("image_history_length", 0), "image_history_interval": data_config.get("image_history_interval", 1), "future_image_length": data_config.get("future_image_length", 0), "future_image_interval": data_config.get("future_image_interval", 1), "future_image_indices": data_config.get("future_image_indices", None), "max_pixels": data_config.get("max_pixels", 1280 * 28 * 28), "min_pixels": data_config.get("min_pixels", 4 * 28 * 28), "image_factor": data_config.get("image_factor", 28), } ) # 3) Action and time series params.update( { "predict_action_keys": data_config.get("predict_action_keys", []), "obs_action_keys": data_config.get("obs_action_keys", []), "action_horizon": data_config.get("action_horizon", 0), "action_history_length": data_config.get("action_history_length", 0), "action_horizon_flow": data_config.get( "action_horizon_flow", yaml_dict.get("action_horizon_flow", 32) ), "action_horizon_ar": data_config.get("action_horizon_ar", 0), "left_padding": data_config.get("left_padding", True), "right_padding": data_config.get("right_padding", True), "dof_config": yaml_dict.get( "dof_config", data_config.get("dof_config", {}) ), "agent_pos_config": yaml_dict.get( "agent_pos_config", data_config.get("agent_pos_config", {}) ), "state_augmentation_prob": data_config.get( "state_augmentation_prob", 0.05 ), "state_drop_prob": data_config.get("state_drop_prob", 0.0), } ) # 4) Instruction and multimodal params.update( { "default_instruction": data_config.get("default_instruction", ""), "instruction_path": data_config.get("instruction_path", None), "instruction_key": data_config.get("instruction_key", None), "multimodal_chunk_size": data_config.get("multimodal_chunk_size", 500), "generate_subtask_ratio": data_config.get( "generate_subtask_ratio", 0.0 ), "cot_ratio": data_config.get("cot_ratio", 0.0), "multimodal_data_ratio": data_config.get("multimodal_data_ratio", 0.25), "instruction_key_prob": data_config.get("instruction_key_prob", None), "trunc_action_with_instruction": data_config.get( "trunc_action_with_instruction", True ), "use_embodied_system_prompt_ratio": data_config.get( "use_embodied_system_prompt_ratio", yaml_dict.get("use_embodied_system_prompt_ratio", 0.0), ), } ) # 5) Data cleaning and alignment (validation/augmentation/framework constraints) params.update( { "filter_angle_outliers": data_config.get( "filter_angle_outliers", False ), "trim_stationary": data_config.get("trim_stationary", False), "use_state_string_representation": data_config.get( "use_state_string_representation", yaml_dict.get("use_state_string_representation", False), ), "pad_prefix_to_same_length": data_config.get( "pad_prefix_to_same_length", False ), "put_ar_predict_in_postfix": data_config.get( "put_ar_predict_in_postfix", False ), # "pad_to_128_multiple": data_config.get("pad_to_128_multiple", True), "padding_side": data_config.get("padding_side", "left"), "max_seqlen": yaml_dict.get("max_seqlen", 768), "model_type": yaml_dict.get("model_type", "qwen2_5"), "model_config_path": yaml_dict.get("qwen_vl_act_config_path", None), "low_dim_obs_horizon": data_config.get("low_dim_obs_horizon", 1), } ) # Only keep valid fields defined in dataclass valid_fields = {f.name for f in cls.__dataclass_fields__.values()} filtered = {k: v for k, v in params.items() if k in valid_fields} return cls(**filtered) # ---------------------------------------------------------------------- # Dict-style access (for compatibility with existing calls) # ---------------------------------------------------------------------- def __getitem__(self, key: str): try: return getattr(self, key) except AttributeError: raise KeyError(f"'{key}' not found in {self.__class__.__name__}") def __setitem__(self, key: str, value): setattr(self, key, value) def __contains__(self, key: str) -> bool: return hasattr(self, key) def keys(self): return self.__dict__.keys() def values(self): return self.__dict__.values() def items(self): return self.__dict__.items() class InferConfig: def __init__( self, checkpoint_path: str | None = None, train_config_path: str | None = None, robot_host: str = "0.0.0.0", robot_port: int = 33723, robot_id: str = "10053", robot_type: str = "desktop", # ["desktop", "turtle"] robot_action_start_ratio: float = 0, # Action execution start ratio robot_action_end_ratio: float = 0.8, # Action execution end ratio robot_action_interpolate_multiplier: int = 70, # Action interpolation robot_use_joint_angle_control: bool = False, # Use joint control (model must be joint prediction model) turtle_as_desktop: bool = False, # Use turtle body for desktop operation, fixed chassis head movement, head camera, and chassis height action_horizon: int = 10, # Please correctly fill in the model's horizon action_dim: int | None = None, model_device: str = "cuda:0", num_inference_timesteps: int = 10, norm_key: str = "x2_normal", cam_names: list[str] = ["face_view", "right_wrist_view"], ): # Private attribute for storing path assert checkpoint_path is not None self._checkpoint_path = checkpoint_path if os.path.exists(os.path.join(checkpoint_path, "normalizer_action.pth")): self.normalizer_action_path = os.path.join( checkpoint_path, "normalizer_action.pth" ) if os.path.exists(os.path.join(checkpoint_path, "normalizer_propri.pth")): self.normalizer_propri_path = os.path.join( checkpoint_path, "normalizer_propri.pth" ) self.model_path = checkpoint_path self.action_tokenizer_path = "/x2robot_v2/Models/fast/" # Other configuration attributes self.robot_host = robot_host self.robot_port = robot_port self.robot_type = robot_type # ["desktop", "turtle"] self.robot_id = robot_id self.robot_action_start_ratio = robot_action_start_ratio self.robot_action_end_ratio = robot_action_end_ratio self.robot_action_interpolate_multiplier = robot_action_interpolate_multiplier self.robot_use_joint_angle_control = ( robot_use_joint_angle_control # Use joint angle control ) self.turtle_as_desktop = turtle_as_desktop self._action_horizon = ( action_horizon # Default controlled by train config's flow action horizon ) self._action_dim = action_dim # Default determined by train config's dof config self.action_dim = action_dim self.pred_horizon = action_horizon self.predict_mode = "diffusion" self.camera_key = cam_names self.model_device = model_device self.num_inference_timesteps = ( num_inference_timesteps # flow matching related config ) # Initialize config objects self.train_config: dict = {} self.model_config = None self.data_config = None self.norm_key = norm_key self.cam_names = cam_names # Load all configs self._load_all_configs(train_config_path) @property def checkpoint_path(self) -> str | None: return self._checkpoint_path @checkpoint_path.setter def checkpoint_path(self, value: str | None): """When checkpoint_path is updated, reload all configs""" if self._checkpoint_path != value: self._checkpoint_path = value self._load_all_configs() @property def action_horizon(self) -> int: return self._action_horizon @action_horizon.setter def action_horizon(self, value: int): self._action_horizon = value @property def action_dim(self) -> int | None: return self._action_dim @action_dim.setter def action_dim(self, value: int | None): self._action_dim = value def _load_all_configs(self, train_config_path=None): """Unified entry point for loading all configs""" self._load_train_config(train_config_path) self._load_model_config() self._load_data_config() # Update action_horizon and action_dim (if needed) if self._action_horizon is None: self._action_horizon = self.train_config.get("data", {}).get( "action_horizon_flow", 32 ) assert self._action_horizon is not None and self._action_horizon > 0 if self._action_dim is None: self._action_dim = sum(self.train_config.get("dof_config", {}).values()) def _load_train_config(self, train_config_path): if train_config_path is None: train_config_path = os.path.join(self._checkpoint_path, "config.yml") with open(train_config_path, "r") as f: self.train_config = yaml.load(f, Loader=yaml.FullLoader) ckpt_dir = self._checkpoint_path preprocessor_file = os.path.join(ckpt_dir, "preprocessor_config.json") if os.path.exists(preprocessor_file): print(f"[LoadConfig] Found {preprocessor_file}, override processor_path.") self.train_config["processor_path"] = ckpt_dir tokenizer_file = os.path.join(ckpt_dir, "tokenizer.json") tokenizer_config_file = os.path.join(ckpt_dir, "tokenizer_config.json") if "action_tokenizer_path" in self.train_config and not os.path.exists( self.train_config["action_tokenizer_path"] ): if os.path.exists(tokenizer_file) and os.path.exists(tokenizer_config_file): print( f"[LoadConfig] Found tokenizer files in {ckpt_dir}, override action_tokenizer_path." ) self.train_config["action_tokenizer_path"] = ckpt_dir else: print("[LoadConfig] Cannot load action tokenizer! ") def _load_model_config(self): ckpt_config_path = os.path.join(self._checkpoint_path, "config.json") resolved_cfg_path = None if os.path.exists(ckpt_config_path): # Prefer checkpoint config resolved_cfg_path = ckpt_config_path print(f"[LoadModelConfig] Using checkpoint config.json: {ckpt_config_path}") else: # Fallback to original config path fallback_cfg = self.train_config.get("qwen_vl_act_config_path", None) if fallback_cfg is not None: resolved_cfg_path = fallback_cfg print(f"[LoadModelConfig] Using fallback act config: {fallback_cfg}") if resolved_cfg_path is None or (not os.path.exists(resolved_cfg_path)): raise ValueError( f"[LoadModelConfig] Cannot load model config! " f"Checked:\n" f" - Checkpoint config.json: {ckpt_config_path}\n" f" - Fallback path: {self.train_config.get('qwen_vl_act_config_path', None)}" ) # Save back to config for consistency self.train_config["qwen_vl_act_config_path"] = resolved_cfg_path model_type = self.train_config["model_type"] if model_type == "qwen2_5": from wall_x.model.qwen2_5_based import Qwen2_5_VLConfig ConfigClass = Qwen2_5_VLConfig # elif model_type == "qwen3": # from wall_x.model.qwen3_based import Qwen3VLConfig # ConfigClass = Qwen3VLConfig else: raise ValueError(f"[LoadModelConfig] Unsupported model type: {model_type}") print(f"[LoadModelConfig] Loading model config from: {resolved_cfg_path}") self.model_config = ConfigClass.from_pretrained(resolved_cfg_path) self.model_config = update_model_config(self.train_config, self.model_config) self.model_config._attn_implementation = "sdpa" self.model_config.vision_config._attn_implementation = "flash_attention_2" print("[LoadModelConfig] Model config loaded and updated successfully.") def _load_data_config(self): self.data_config = X2RDataConfig.from_yaml_dict(self.train_config) if __name__ == "__main__": config = InferConfig() print(config.train_config)