Refresh OSS export from refactor_2604: fused CUDA kernels, vanilla we…

THIRD_PARTY_NOTICES.md

@@ -0,0 +1,40 @@
+# Third-Party Notices
+
+This repository includes or adapts code from the projects listed below. Source
+files that already carry upstream copyright and license headers keep those
+headers in place.
+
+## Apache-2.0 Components
+
+The full Apache License 2.0 text is provided in `LICENSE`.
+
+- `scripts/normalize.py` adapts code or structure from
+  `Physical-Intelligence/openpi`, licensed under Apache-2.0.
+- `wall_x/model/qact/qwen2_5/configuration_qwen2_5_vl.py` and
+  `wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl.py` include Qwen /
+  Hugging Face derived code, licensed under Apache-2.0.
+
+## MIT Components
+
+LIBERO evaluation helpers in `wall_x/_vendor/harrix/envs/libero*.py` interface
+with LIBERO components, licensed under the MIT License:
+
+Copyright (c) 2023 Lifelong Robot Learning
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

requirements.txt

@@ -24,3 +24,8 @@ tqdm>=4.67,<5
 tyro>=0.9.0,<1.0
 wandb>=0.23,<0.25
 matplotlib>=3.10,<4
+imageio>=2.34,<3
+imageio-ffmpeg>=0.5,<0.7
+# DMuon optimizer (recommended for LIBERO finetuning; see the *_dmuon.yml
+# configs). Public, installed from git (not published on PyPI).
+dmuon @ git+https://github.com/X-Square-Robot/dmuon.git

scripts/draw_openloop_plot.py

@@ -7,16 +7,13 @@
 
 Reference: ``infer_openloop_websocket.py`` (websocket client + open-loop loop).
 
-For single-arm LIBERO checkpoints, start the server with raw model output::
-
-    python -m wall_x._vendor.harrix.serving.launch_serving \\
-      --env X2ROBOT --port 32194 \\
-      --no-serialize-actions \\
-      model-config:server-model-config \\
-      --model-config.checkpoint-path /path/to/ckpt \\
-      --model-config.train-config-path /path/to/libero.yml \\
-      --model-config.action-horizon 10 \\
-      --model-config.robot-type desktop
+For single-arm LIBERO checkpoints, start the vanilla server (raw model output)::
+
+    bash scripts/run_serving.sh \\
+      --checkpoint-path /path/to/ckpt \\
+      --train-config-path /path/to/libero.yml \\
+      --action-horizon 10 \\
+      --port 32194
 """
 
 from __future__ import annotations
@@ -361,15 +358,9 @@ def decode_predict_action(
         if parts:
             return np.concatenate(parts, axis=1).astype(np.float32)
 
-    try:
-        from wall_x._vendor.harrix.envs.libero_common import decode_chunk
-
-        chunk = decode_chunk(pa, train_config)
-        if chunk.ndim == 2 and chunk.shape[1] > 0:
-            return chunk.astype(np.float32)
-    except Exception:
-        pass
-
+    # No dof_config available: fall back to the raw model action rows. (The
+    # primary path above already drops the virtual action_padding tail, which is
+    # the OSS-correct decode.)
     return pa.astype(np.float32)
 
 

scripts/fake_inference.py

@@ -13,7 +13,6 @@
 from pathlib import Path
 
 import numpy as np
-import torch
 
 
 def _ensure_local_harrix_on_path() -> None:
@@ -56,35 +55,13 @@ def _build_fake_observation(seed: int, image_size: int, gripper_dim: int) -> dic
 
 
 def _predict_raw_action(adapter, payload: dict) -> np.ndarray:
-    from wall_x._vendor.harrix.envs.libero_common import encode_proprio
+    """Run one flow-inference pass via the adapter's batched predict path.
 
-    observation = encode_proprio(
-        payload["observation"], adapter._train_config, adapter._action_horizon
-    )
-    prefix, postfix = adapter._get_flow_prompt(payload["instruction"])
-    batch_inputs = adapter._construct_model_input([observation], [prefix], [postfix])
-
-    dataset_names = batch_inputs["dataset_names"]
-    padding = torch.zeros(
-        (len(dataset_names), 1, adapter._action_dim),
-        dtype=torch.float32,
-    )
-    padding_action = adapter._normalizer_action.normalize_data(
-        padding, dataset_names
-    ).to(batch_inputs["input_ids"].device)
-
-    model_output = adapter._model.generate_flow_action(
-        action_horizon=adapter._action_horizon,
-        action_dim=adapter._action_dim,
-        num_inference_timesteps=adapter._num_inference_timesteps,
-        padding_action=padding_action,
-        noise=payload.get("noise"),
-        **batch_inputs,
-    )
-    predict_action = model_output["predict_action"]
-    if isinstance(predict_action, torch.Tensor):
-        predict_action = predict_action.detach().cpu().numpy()
-    return np.asarray(predict_action[0])
+    ``predict_batch`` is the same path the LIBERO evaluation uses; it handles
+    proprio encoding, prompt construction, and flow sampling internally.
+    """
+    actions = adapter.predict_batch([payload])
+    return np.asarray(actions[0])
 
 
 def parse_args() -> argparse.Namespace:
@@ -117,9 +94,9 @@ def main() -> int:
     args = parse_args()
     _ensure_local_harrix_on_path()
 
-    import wall_x._vendor.harrix.adapters  # noqa: F401  register model adapters
-    from wall_x._vendor.harrix.adapters.registry import build_adapter
-    from wall_x._vendor.harrix.eval_config import (
+    import wall_x._vendor.harrix.e2e_infer.adapters  # noqa: F401  register model adapters
+    from wall_x._vendor.harrix.e2e_infer.adapters.registry import build_adapter
+    from wall_x._vendor.harrix.e2e_infer.config import (
         EvalConfig,
         LiberoEnvParams,
         autofill_from_checkpoint,
@@ -137,7 +114,9 @@ def main() -> int:
     cfg = autofill_from_checkpoint(cfg)
 
     adapter = build_adapter(cfg)
-    gripper_dim = _right_gripper_dim_from_config(getattr(adapter, "_train_config", {}))
+    gripper_dim = _right_gripper_dim_from_config(
+        {"agent_pos_config": getattr(adapter._typed_cfg.task, "agent_pos_config", {})}
+    )
     payload = {
         "observation": _build_fake_observation(args.seed, args.image_size, gripper_dim),
         "instruction": args.instruction,

scripts/infer_libero.py

@@ -15,6 +15,7 @@
 
 import yaml
 
+
 LIBERO_DEFAULT_MAX_INFER_TIMES = {
     "libero_spatial": 22,
     "libero_object": 28,
@@ -211,10 +212,7 @@ def main() -> int:
     args = parse_args()
     _ensure_local_harrix_on_path()
 
-    from wall_x._vendor.harrix.eval_config import (
-        autofill_from_checkpoint,
-        load_eval_config,
-    )
+    from wall_x._vendor.harrix.e2e_infer.config import autofill_from_checkpoint, load_eval_config
 
     raw = _load_or_build_raw_config(args)
 
@@ -225,7 +223,7 @@ def main() -> int:
     cfg = autofill_from_checkpoint(load_eval_config(tmp_config))
     if cfg.runtime.driver_mode != "in_process":
         raise ValueError("Only driver_mode='in_process' is supported")
-    from wall_x._vendor.harrix.drivers.inproc import run
+    from wall_x._vendor.harrix.e2e_infer.drivers.inproc import run
 
     run(cfg)
     return 0

scripts/normalize.py

@@ -0,0 +1,163 @@
+# Adapted from Physical-Intelligence/openpi:
+# https://github.com/Physical-Intelligence/openpi/blob/650c5b0283a49c42784fb5055a0507da2c6d347d/src/openpi/shared/normalize.py
+# Upstream license: Apache-2.0.
+import json
+import pathlib
+
+import numpy as np
+import numpydantic
+import pydantic
+
+
+@pydantic.dataclasses.dataclass
+class NormStats:
+    mean: numpydantic.NDArray
+    std: numpydantic.NDArray
+    q01: numpydantic.NDArray | None = None  # 1st quantile
+    q99: numpydantic.NDArray | None = None  # 99th quantile
+
+
+class RunningStats:
+    """Compute running statistics of a batch of vectors."""
+
+    def __init__(self):
+        self._count = 0
+        self._mean = None
+        self._mean_of_squares = None
+        self._min = None
+        self._max = None
+        self._histograms = None
+        self._bin_edges = None
+        self._num_quantile_bins = 5000  # for computing quantiles on the fly
+
+    def update(self, batch: np.ndarray) -> None:
+        """
+        Update the running statistics with a batch of vectors.
+
+        Args:
+            vectors (np.ndarray): An array where all dimensions except the last are batch dimensions.
+        """
+        batch = batch.reshape(-1, batch.shape[-1])
+        num_elements, vector_length = batch.shape
+        if self._count == 0:
+            self._mean = np.mean(batch, axis=0)
+            self._mean_of_squares = np.mean(batch**2, axis=0)
+            self._min = np.min(batch, axis=0)
+            self._max = np.max(batch, axis=0)
+            self._histograms = [
+                np.zeros(self._num_quantile_bins) for _ in range(vector_length)
+            ]
+            self._bin_edges = [
+                np.linspace(
+                    self._min[i] - 1e-10,
+                    self._max[i] + 1e-10,
+                    self._num_quantile_bins + 1,
+                )
+                for i in range(vector_length)
+            ]
+        else:
+            if vector_length != self._mean.size:
+                raise ValueError(
+                    "The length of new vectors does not match the initialized vector length."
+                )
+            new_max = np.max(batch, axis=0)
+            new_min = np.min(batch, axis=0)
+            max_changed = np.any(new_max > self._max)
+            min_changed = np.any(new_min < self._min)
+            self._max = np.maximum(self._max, new_max)
+            self._min = np.minimum(self._min, new_min)
+
+            if max_changed or min_changed:
+                self._adjust_histograms()
+
+        self._count += num_elements
+
+        batch_mean = np.mean(batch, axis=0)
+        batch_mean_of_squares = np.mean(batch**2, axis=0)
+
+        # Update running mean and mean of squares.
+        self._mean += (batch_mean - self._mean) * (num_elements / self._count)
+        self._mean_of_squares += (batch_mean_of_squares - self._mean_of_squares) * (
+            num_elements / self._count
+        )
+
+        self._update_histograms(batch)
+
+    def get_statistics(self) -> NormStats:
+        """
+        Compute and return the statistics of the vectors processed so far.
+
+        Returns:
+            dict: A dictionary containing the computed statistics.
+        """
+        if self._count < 2:
+            raise ValueError("Cannot compute statistics for less than 2 vectors.")
+
+        variance = self._mean_of_squares - self._mean**2
+        stddev = np.sqrt(np.maximum(0, variance))
+        q01, q99 = self._compute_quantiles([0.01, 0.99])
+        return NormStats(mean=self._mean, std=stddev, q01=q01, q99=q99)
+
+    def _adjust_histograms(self):
+        """Adjust histograms when min or max changes."""
+        for i in range(len(self._histograms)):
+            old_edges = self._bin_edges[i]
+            new_edges = np.linspace(
+                self._min[i], self._max[i], self._num_quantile_bins + 1
+            )
+
+            # Redistribute the existing histogram counts to the new bins
+            new_hist, _ = np.histogram(
+                old_edges[:-1], bins=new_edges, weights=self._histograms[i]
+            )
+
+            self._histograms[i] = new_hist
+            self._bin_edges[i] = new_edges
+
+    def _update_histograms(self, batch: np.ndarray) -> None:
+        """Update histograms with new vectors."""
+        for i in range(batch.shape[1]):
+            hist, _ = np.histogram(batch[:, i], bins=self._bin_edges[i])
+            self._histograms[i] += hist
+
+    def _compute_quantiles(self, quantiles):
+        """Compute quantiles based on histograms."""
+        results = []
+        for q in quantiles:
+            target_count = q * self._count
+            q_values = []
+            for hist, edges in zip(self._histograms, self._bin_edges, strict=True):
+                cumsum = np.cumsum(hist)
+                idx = np.searchsorted(cumsum, target_count)
+                q_values.append(edges[idx])
+            results.append(np.array(q_values))
+        return results
+
+
+class _NormStatsDict(pydantic.BaseModel):
+    norm_stats: dict[str, NormStats]
+
+
+def serialize_json(norm_stats: dict[str, NormStats]) -> str:
+    """Serialize the running statistics to a JSON string."""
+    return _NormStatsDict(norm_stats=norm_stats).model_dump_json(indent=2)
+
+
+def deserialize_json(data: str) -> dict[str, NormStats]:
+    """Deserialize the running statistics from a JSON string."""
+    return _NormStatsDict(**json.loads(data)).norm_stats
+
+
+def save(directory: pathlib.Path | str, norm_stats: dict[str, NormStats]) -> None:
+    """Save the normalization stats to a directory."""
+    path = pathlib.Path(directory) / "norm_stats.json"
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(serialize_json(norm_stats))
+
+
+def load(directory: pathlib.Path | str) -> dict[str, NormStats]:
+    """Load the normalization stats from a directory."""
+    path = pathlib.Path(directory) / "norm_stats.json"
+    if not path.exists():
+        raise FileNotFoundError(f"Norm stats file not found at: {path}")
+    return deserialize_json(path.read_text())