Add Wall-X serving and Turtle2 TCP WebSocket bridge

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