Add Wall-X serving and Turtle2 TCP WebSocket bridge

wall_x/_vendor/harrix/utils/ckpt_load.py

@@ -7,6 +7,8 @@ are intentionally left to the adapter.
 from __future__ import annotations
 
 import os
+import json
+import math
 from typing import Callable, Optional
 
 import torch
@@ -17,6 +19,58 @@ def _noop_log(_msg: str, **_kw) -> None:
     pass
 
 
+def resolve_lora_scale(
+    checkpoint_path: str,
+    train_config: dict,
+    state_dict: dict,
+    log_fn: Optional[Callable] = None,
+) -> float:
+    """Read PEFT alpha/r from an exported config and validate it against weights."""
+    log_fn = log_fn or _noop_log
+    ranks = {
+        int(tensor.shape[0])
+        for name, tensor in state_dict.items()
+        if name.endswith(".lora_A.default.weight")
+    }
+    if not ranks:
+        return 1.0
+    if len(ranks) != 1:
+        raise ValueError(f"LoRA weights contain multiple ranks: {sorted(ranks)}")
+    rank = ranks.pop()
+
+    candidates = [
+        os.path.join(checkpoint_path, "lora_config.json"),
+        os.path.join(checkpoint_path, "adapter_config.json"),
+        (train_config.get("model") or {}).get("lora_config_path"),
+    ]
+    source = next((path for path in candidates if path and os.path.isfile(path)), None)
+    if source is None:
+        log_fn(
+            "LoRA scale metadata is missing; using legacy alpha/r=2.0. "
+            "Export the training LoRA JSON as checkpoint/lora_config.json "
+            "before deploying a checkpoint with changed LoRA settings."
+        )
+        return 2.0
+
+    with open(source, encoding="utf-8") as stream:
+        config = json.load(stream)
+    configured_rank = int(config.get("lora_r", config.get("r", rank)))
+    if configured_rank != rank:
+        raise ValueError(
+            f"LoRA rank mismatch: {source} has {configured_rank}, weights have {rank}"
+        )
+    if config.get("rank_pattern") or config.get("alpha_pattern"):
+        raise ValueError(f"Per-layer LoRA scaling in {source} is unsupported")
+    alpha = config.get("lora_alpha")
+    if alpha is None:
+        raise ValueError(f"LoRA config {source} lacks lora_alpha")
+    scale = float(alpha) / (math.sqrt(rank) if config.get("use_rslora") else rank)
+    if not math.isfinite(scale) or scale <= 0:
+        raise ValueError(f"Invalid LoRA scale {scale} from {source}")
+    log_fn(f"LoRA scale={scale:g} from {source} (alpha={alpha}, rank={rank})")
+    return scale
+
+
 def _align_checkpoint_tensor(
     param: torch.Tensor,
     target: torch.Tensor,
@@ -64,22 +118,58 @@ def _align_checkpoint_tensor(
 
 
 def reshape_compatible_state_dict(
-    state_dict: dict, model_sd: dict, log_fn: Optional[Callable] = None
+    state_dict: dict,
+    model_sd: dict,
+    log_fn: Optional[Callable] = None,
+    lora_scale: float = 2.0,
 ) -> dict:
     """Align checkpoint tensors to the target model shapes via crop / pad."""
     log_fn = log_fn or _noop_log
+    # Training checkpoints keep PEFT modules (base_layer/lora_A/lora_B),
+    # while inference uses the plain fused projection tensors. Materialize
+    # those tensors here so serving does not silently drop the LoRA update.
+    merged = dict(state_dict)
+    for name, base in list(state_dict.items()):
+        suffix = ".base_layer.weight"
+        if not name.endswith(suffix):
+            continue
+        stem = name[: -len(suffix)]
+        a_name = stem + ".lora_A.default.weight"
+        b_name = stem + ".lora_B.default.weight"
+        if a_name not in state_dict or b_name not in state_dict:
+            continue
+        a = state_dict[a_name].to(dtype=torch.float32)
+        b = state_dict[b_name].to(dtype=torch.float32)
+        merged[name.replace(".base_layer.weight", ".weight")] = (
+            base.to(dtype=torch.float32) + lora_scale * (b @ a)
+        ).to(dtype=base.dtype)
+        del merged[name]
+        del merged[a_name]
+        del merged[b_name]
+        log_fn(f"Merged LoRA weights for {stem}")
+    state_dict = merged
     out = {}
     for name, param in state_dict.items():
-        if name not in model_sd:
+        # Training with PEFT wraps the backbone under ``base_model.model``;
+        # serving instantiates the unwrapped model. Normalize that prefix so
+        # fine-tuned backbone weights are applied instead of being ignored.
+        target_name = name
+        if ".base_layer." in target_name:
+            target_name = target_name.replace(".base_layer.", ".")
+        if target_name not in model_sd and target_name.startswith(
+            "model.base_model.model."
+        ):
+            target_name = "model." + target_name[len("model.base_model.model.") :]
+        if target_name not in model_sd:
             log_fn(f"Not used parameter: {name}")
             continue
-        target = model_sd[name]
+        target = model_sd[target_name]
         if param.shape == target.shape:
-            out[name] = param
+            out[target_name] = param
             continue
-        aligned = _align_checkpoint_tensor(param, target, name, log_fn)
+        aligned = _align_checkpoint_tensor(param, target, target_name, log_fn)
         if aligned is not None:
-            out[name] = aligned
+            out[target_name] = aligned
     return out
 
 

wall_x/_vendor/harrix/utils/normalizer.py

@@ -1,8 +1,10 @@
 """Build action/proprio normalizers and resolve the effective norm key.
 
 Public inference artifacts must carry their own normalization data. This module
-uses checkpoint-local ``norm_stats.json`` first, then checkpoint-side normalizer
-state dicts, and finally an explicit ``customized_action_statistic_dof`` path.
+uses checkpoint-local ``norm_stats.json`` for the 7D Euler layout, and uses
+checkpoint-side normalizer state dicts for checkpoints whose train config contains
+6D rotation fields. It finally supports an explicit
+``customized_action_statistic_dof`` path.
 It does not fall back to internal default action statistics.
 """
 
@@ -59,6 +61,19 @@ def _missing_normalizer_error(checkpoint_path: str, train_config: dict) -> FileN
     )
 
 
+def _uses_rotation_6d(train_config: dict) -> bool:
+    """Return whether the checkpoint was trained with 6D rotation fields."""
+    for layout_name in ("dof_config", "agent_pos_config", "ar_dof_config"):
+        layout = train_config.get(layout_name)
+        if layout is None and isinstance(train_config.get("task"), dict):
+            layout = train_config["task"].get(layout_name)
+        if isinstance(layout, dict) and any(
+            "rotation_6d" in str(key).lower() for key in layout
+        ):
+            return True
+    return False
+
+
 def build_normalizers(
     checkpoint_path: str,
     train_config: dict,
@@ -66,14 +81,24 @@ def build_normalizers(
 ) -> tuple[Normalizer, Normalizer, str]:
     """Return action/proprio normalizers and the resolved norm key."""
     norm_stats_path = os.path.join(checkpoint_path, "norm_stats.json")
-    if os.path.exists(norm_stats_path):
+    action_pth = os.path.join(checkpoint_path, "normalizer_action.pth")
+    propri_pth = os.path.join(checkpoint_path, "normalizer_propri.pth")
+    prefer_checkpoint_pth = (
+        os.path.exists(action_pth)
+        and os.path.exists(propri_pth)
+        and _uses_rotation_6d(train_config)
+    )
+    if os.path.exists(norm_stats_path) and not prefer_checkpoint_pth:
         propri_stats = _load_norm_stats(norm_stats_path, "observation.state")
         action_stats = _load_norm_stats(norm_stats_path, "action")
         normalizer_propri = Normalizer.from_lerobot_norm_stats(propri_stats, norm_key)
         normalizer_action = Normalizer.from_lerobot_norm_stats(action_stats, norm_key)
     else:
-        action_pth = os.path.join(checkpoint_path, "normalizer_action.pth")
-        propri_pth = os.path.join(checkpoint_path, "normalizer_propri.pth")
+        if prefer_checkpoint_pth:
+            logger.info(
+                "Using checkpoint normalizer .pth files because train config "
+                "contains rotation_6D; ignoring 7D norm_stats.json"
+            )
         custom_stats = _load_custom_action_stats(train_config)
         if custom_stats is None and (not os.path.exists(action_pth) or not os.path.exists(propri_pth)):
             raise _missing_normalizer_error(checkpoint_path, train_config)