Add Wall-X serving and Turtle2 TCP WebSocket bridge

wall_x/trainer/adapters/base_adapter.py

@@ -94,6 +94,10 @@ class ModelAdapter(ABC):
         """Load pretrained weights, resize embeddings, set normalizers, etc."""
         ...
 
+    def finalize_model_after_weight_load(self, model, model_config):
+        """Apply wrappers that must be created after loading base model weights."""
+        return model
+
     # ---- FSDP wrapping ----
     @abstractmethod
     def get_transformer_layer_cls(self):

wall_x/trainer/adapters/vla_model_adapter.py

@@ -1,3 +1,4 @@
+import json
 import os
 
 import torch
@@ -107,6 +108,8 @@ class VLAdapter(ModelAdapter):
         flat = dataclasses.asdict(self.cfg.model)
         flat["model_type"] = self.cfg.model_type
         flat["data"] = dict(self.cfg._raw_data or {})
+        flat["data"]["action_horizon"] = self.cfg.task.action_horizon
+        flat["data"]["action_horizon_flow"] = self.cfg.task.action_horizon_flow
         flat["dof_config"] = self.cfg.task.dof_config
         flat["agent_pos_config"] = self.cfg.task.agent_pos_config
         if self.cfg.task.ar_dof_config is not None:
@@ -145,6 +148,16 @@ class VLAdapter(ModelAdapter):
         else:
             model_config = ConfigClass.from_pretrained(qwen_vl_act_config_path)
 
+        lora_config_path = getattr(self.cfg.model, "lora_config_path", None)
+        if lora_config_path:
+            with open(lora_config_path, encoding="utf-8") as f:
+                lora_overrides = json.load(f)
+            if not isinstance(lora_overrides, dict):
+                raise ValueError("model.lora_config_path must contain a JSON object")
+            for key, value in lora_overrides.items():
+                setattr(model_config, key, value)
+            self.logger.info("Loaded LoRA overrides from %s", lora_config_path)
+
         assert self.model_type in model_config.model_type, (
             f"Mismatch of model type: model type in config file is "
             f"{model_config.model_type}, but the model type is {self.model_type}."
@@ -156,12 +169,36 @@ class VLAdapter(ModelAdapter):
     def create_model(self, processor, tokenizer_mixin, model_config):
         ModelClass, _ = self._get_model_and_config_class()
         use_selective_recompute = self.cfg.distributed.use_selective_recompute
-        return ModelClass(
-            model_config,
-            processor,
-            tokenizer_mixin=tokenizer_mixin,
-            use_selective_recompute=use_selective_recompute,
+        # PEFT changes parameter names (base_model/base_layer). Build the plain
+        # model first so Qwen and Wall-OSS checkpoints load against native keys.
+        use_lora = bool(getattr(model_config, "use_lora", False))
+        if use_lora:
+            model_config.use_lora = False
+        try:
+            return ModelClass(
+                model_config,
+                processor,
+                tokenizer_mixin=tokenizer_mixin,
+                use_selective_recompute=use_selective_recompute,
+            )
+        finally:
+            if use_lora:
+                model_config.use_lora = True
+
+    def finalize_model_after_weight_load(self, model, model_config):
+        """Inject LoRA only after all unwrapped checkpoint weights are loaded."""
+        if not getattr(model_config, "use_lora", False):
+            return model
+        model.add_lora(
+            r=model_config.lora_r,
+            lora_alpha=model_config.lora_alpha,
+            target_modules=model_config.lora_target_modules,
+            lora_dropout=model_config.lora_dropout,
         )
+        if getattr(model_config, "lora_train_action_expert", False):
+            model._set_lora_action_expert_trainable()
+        self.logger.info("Applied LoRA after loading base and Wall-OSS weights")
+        return model
 
     def load_weights(self, model, normalizer_action, normalizer_propri, **kwargs):
         import copy

wall_x/trainer/fsdp_trainer/fsdp_trainer.py

@@ -95,6 +95,18 @@ class FSDPTrainer(DistributedTrainer):
         self.load_processor()
         self.load_model()
 
+        # Single-file pretrained checkpoints use native model parameter names.
+        # Load them before PEFT or other parameter-name-changing wrappers.
+        if self._resume_from_single_file():
+            self.load_state_dict(
+                self.model,
+                {"ckpt": self.cfg.checkpoint.resume_from},
+            )
+
+        self.model = self.adapter.finalize_model_after_weight_load(
+            self.model, self.model_config
+        )
+
         # DDP needs to see correct requires_grad at wrap time, so freeze first.
         self._freeze_params_if_needed(self.model)
 
@@ -107,12 +119,6 @@ class FSDPTrainer(DistributedTrainer):
             name: p.shape for name, p in self.model.named_parameters()
         }
 
-        if self._resume_from_single_file():
-            self.load_state_dict(
-                self.model,
-                {"ckpt": self.cfg.checkpoint.resume_from},
-            )
-
         self._wrap_model(self.model)
         self._create_optimizer()
         self._create_scheduler()

wall_x/trainer/trainer_utils.py

@@ -161,6 +161,8 @@ def load_wallx_processors_from_cfg(
     flat = dataclasses.asdict(cfg.model)
     flat["model_type"] = cfg.model_type
     flat["data"] = dict(cfg._raw_data or {})
+    flat["data"]["action_horizon"] = cfg.task.action_horizon
+    flat["data"]["action_horizon_flow"] = cfg.task.action_horizon_flow
     flat["dof_config"] = cfg.task.dof_config
     flat["agent_pos_config"] = cfg.task.agent_pos_config
     if cfg.task.ar_dof_config is not None: