Update Wall-X to 1.1.0 (#104)
This commit is contained in:
Starrick Liu
@@ -0,0 +1,142 @@
|
||||
# Scripts
|
||||
|
||||
This directory contains the public Wall-X command-line helpers. Run the examples
|
||||
below from the repository root, using `python scripts/...` and `bash scripts/...`.
|
||||
Pass file and directory paths explicitly.
|
||||
|
||||
## Inference smoke test
|
||||
|
||||
Use `fake_inference.py` to verify that a checkpoint can be loaded and can
|
||||
produce one action chunk from a synthetic LIBERO-style observation.
|
||||
|
||||
```bash
|
||||
python scripts/fake_inference.py --checkpoint-path /path/to/checkpoint
|
||||
```
|
||||
|
||||
If the training config is not stored next to the checkpoint as `config.yml` or
|
||||
`config.yaml`, pass it explicitly:
|
||||
|
||||
```bash
|
||||
python scripts/fake_inference.py \
|
||||
--checkpoint-path /path/to/checkpoint \
|
||||
--train-config-path /path/to/config.yml
|
||||
```
|
||||
|
||||
## LIBERO evaluation
|
||||
|
||||
`run_libero.sh` is a small shell wrapper around `infer_libero.py`. It requires
|
||||
the optional LIBERO simulator stack:
|
||||
|
||||
```bash
|
||||
pip install -r requirements-libero.txt
|
||||
mkdir -p third_party
|
||||
git clone https://github.com/Lifelong-Robot-Learning/LIBERO.git third_party/LIBERO
|
||||
```
|
||||
|
||||
The launcher checks for LIBERO, robosuite, MuJoCo, PyOpenGL, BDDL, Gym, and
|
||||
h5py before loading the model. If LIBERO is cloned elsewhere, pass
|
||||
`LIBERO_PATH=/path/to/LIBERO`.
|
||||
|
||||
```bash
|
||||
bash scripts/run_libero.sh /path/to/checkpoint
|
||||
```
|
||||
|
||||
Useful environment variables:
|
||||
|
||||
```bash
|
||||
CHECKPOINT_PATH=/path/to/checkpoint
|
||||
TRAIN_CONFIG_PATH=/path/to/config.yml
|
||||
TASK_SUITE_NAME=libero_spatial
|
||||
TASK_INDICES=0,1,2
|
||||
NUM_TRIALS_PER_TASK=50
|
||||
CUDA_ID=0
|
||||
SMOKE=1
|
||||
MAX_INFER_TIMES=52
|
||||
```
|
||||
|
||||
`MAX_INFER_TIMES` is optional. When omitted, the launcher uses suite-specific
|
||||
defaults aligned with the LIBERO evaluator: spatial 22, object 28, goal 30,
|
||||
libero_10 52, and libero_90 40 action chunks.
|
||||
|
||||
For full control, call the Python entry directly:
|
||||
|
||||
```bash
|
||||
python scripts/infer_libero.py \
|
||||
--checkpoint-path /path/to/checkpoint \
|
||||
--task-suite-name libero_spatial \
|
||||
--num-trials-per-task 50 \
|
||||
--driver-mode in_process
|
||||
```
|
||||
|
||||
You can also pass a complete eval config:
|
||||
|
||||
```bash
|
||||
python scripts/infer_libero.py --config /path/to/eval_config.yml
|
||||
```
|
||||
|
||||
## WebSocket serving
|
||||
|
||||
`run_serving.sh` launches the Wall-X WebSocket server through the public
|
||||
vendored serving runtime. Pass paths explicitly; the script has no built-in
|
||||
checkpoint path.
|
||||
|
||||
```bash
|
||||
bash scripts/run_serving.sh \
|
||||
--checkpoint-path /path/to/checkpoint \
|
||||
--train-config-path /path/to/config.yml \
|
||||
--port 32195
|
||||
```
|
||||
|
||||
By default the script returns raw model action chunks, which is the expected
|
||||
mode for open-loop plotting. Pass `--serialize-actions` when your client expects
|
||||
robot-serialized actions.
|
||||
|
||||
Useful options:
|
||||
|
||||
```bash
|
||||
CUDA_ID=0
|
||||
ACTION_HORIZON=32
|
||||
IMAGE_PASSING_MODE=base64
|
||||
MAX_BATCH_SIZE=1
|
||||
```
|
||||
|
||||
Additional `launch_serving.py` arguments can be forwarded after `--`:
|
||||
|
||||
```bash
|
||||
bash scripts/run_serving.sh --checkpoint-path /path/to/checkpoint -- \
|
||||
--model-config.norm-key libero_all
|
||||
```
|
||||
|
||||
## Open-loop WebSocket evaluation
|
||||
|
||||
`draw_openloop_plot.py` compares predicted action chunks from a running
|
||||
WebSocket server against LeRobot dataset ground truth. `--dataset-root` and
|
||||
`--train-config` are required and have no built-in default.
|
||||
|
||||
```bash
|
||||
python scripts/draw_openloop_plot.py \
|
||||
--uri ws://127.0.0.1:32195 \
|
||||
--dataset-root /path/to/lerobot_dataset \
|
||||
--train-config /path/to/train_config.yml \
|
||||
--episode-indices 0,1,2 \
|
||||
--save-dir ./openloop_plots
|
||||
```
|
||||
|
||||
## Dataset and checkpoint utilities
|
||||
|
||||
- `compute_norm_stats.py`: compute action normalization statistics for a
|
||||
local LeRobot v3 dataset. The script reads state/action parquet columns
|
||||
directly when available, so image and video columns are not decoded.
|
||||
- `merge_sharded_weights.py`: merge FSDP sharded checkpoint files into a single
|
||||
checkpoint directory.
|
||||
- `merge_tokenizer.py`: merge FAST action tokens into a Qwen2.5-VL processor
|
||||
tokenizer.
|
||||
|
||||
```bash
|
||||
python scripts/merge_tokenizer.py \
|
||||
--processor-path /path/to/Qwen2.5-VL-3B-Instruct \
|
||||
--action-tokenizer-path /path/to/fast_tokenizer \
|
||||
--output-dir /path/to/merged_processor
|
||||
```
|
||||
|
||||
Most scripts support `--help` for their command-line options.
|
||||
+704
-134
@@ -1,182 +1,752 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Compute LeRobot normalization stats (mean, std, q01, q99) for training.
|
||||
|
||||
Writes JSON in the format expected by wall-x training configs::
|
||||
|
||||
{"norm_stats": {
|
||||
"observation.state": {"mean": [...], "std": [...], "q01": [...], "q99": [...]},
|
||||
"action": {"mean": [...], "std": [...], "q01": [...], "q99": [...]}
|
||||
}}
|
||||
|
||||
When ``--train_config`` is provided, the script reads ``data.lerobot_config.repo_id``,
|
||||
``norm_stats_path``, ``task.dof_config``, ``task.agent_pos_config``, and
|
||||
``task.action_horizon`` from the YAML. Per-DOF slices are aggregated separately;
|
||||
keys ending with ``_relative`` use the same relative-pose logic as the LeRobot loader.
|
||||
|
||||
Usage
|
||||
-----
|
||||
Recommended: pass a finetune YAML (paths in the config can be placeholders; override
|
||||
with CLI flags if needed)::
|
||||
|
||||
python scripts/compute_norm_stats.py \\
|
||||
--train_config /path/to/train_config.yml
|
||||
|
||||
Multi-task example::
|
||||
|
||||
python scripts/compute_norm_stats.py \\
|
||||
--train_config /path/to/multitask_config.yml
|
||||
|
||||
Override dataset or output path from the command line::
|
||||
|
||||
python scripts/compute_norm_stats.py \\
|
||||
--train_config /path/to/train_config.yml \\
|
||||
--data_root /path/to/repo_id \\
|
||||
--output_path /path/to/norm_stats_path
|
||||
|
||||
Without a train config (global stats only, no per-DOF relative slices)::
|
||||
|
||||
python scripts/compute_norm_stats.py \\
|
||||
--data_root /path/to/lerobot_dataset \\
|
||||
--output_path /path/to/norm_stats.json
|
||||
|
||||
Requirements
|
||||
------------
|
||||
- Local LeRobot v3 dataset at ``--data_root`` (or ``data.lerobot_config.repo_id``)
|
||||
- ``lerobot>=0.3``, ``datasets``, ``pyarrow``, ``numpy``, ``pyyaml``, ``tqdm``
|
||||
|
||||
After running, set ``norm_stats_path`` in your training YAML to the generated JSON.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Dict, List
|
||||
from tqdm import tqdm
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import yaml
|
||||
from numba import jit, prange
|
||||
from tqdm import tqdm
|
||||
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
SKIP_DOF_KEYS = frozenset(
|
||||
{"velocity_decomposed", "height", "head_actions", "action_padding"}
|
||||
)
|
||||
|
||||
_GEOMETRY_FUNCS_LOADED = False
|
||||
|
||||
|
||||
def write_json(path: Path, data: Dict) -> None:
|
||||
def _ensure_geometry_funcs() -> None:
|
||||
global _GEOMETRY_FUNCS_LOADED
|
||||
global canonicalize_euler_zyx_batch_nb
|
||||
global euler_to_matrix_zyx_batch_nb
|
||||
global matrix_to_euler_zyx_batch_nb
|
||||
global so3_to_matrix_batch_nb
|
||||
|
||||
if _GEOMETRY_FUNCS_LOADED:
|
||||
return
|
||||
|
||||
try:
|
||||
from wall_x._vendor.x2robot_utils.geometry import (
|
||||
canonicalize_euler_zyx_batch_nb,
|
||||
euler_to_matrix_zyx_batch_nb,
|
||||
matrix_to_euler_zyx_batch_nb,
|
||||
so3_to_matrix_batch_nb,
|
||||
)
|
||||
except ImportError as exc:
|
||||
raise RuntimeError(
|
||||
"compute_norm_stats.py requires the vendored x2robot_utils geometry "
|
||||
"helpers. Install Wall-X first so wall_x._vendor is available."
|
||||
) from exc
|
||||
|
||||
_GEOMETRY_FUNCS_LOADED = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class TrainConfigContext:
|
||||
data_root: Path
|
||||
output_path: Path
|
||||
state_key: str
|
||||
action_key: str
|
||||
propri_ranges: dict[str, list[int]]
|
||||
action_ranges: dict[str, list[int]]
|
||||
action_chunk: int
|
||||
dof_config: dict[str, int]
|
||||
agent_pos_config: dict[str, int]
|
||||
|
||||
|
||||
def write_norm_stats(path: Path, norm_stats: dict[str, dict]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(
|
||||
json.dumps(data, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
|
||||
json.dumps({"norm_stats": norm_stats}, ensure_ascii=False, indent=2) + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
def compute_action_statistics(
|
||||
action_data_by_robot: Dict[str, Dict[str, List]]
|
||||
) -> Dict[str, Dict[str, Dict]]:
|
||||
def compute_vector_stats(values: np.ndarray) -> dict[str, list[float]]:
|
||||
if values.ndim == 1:
|
||||
values = values.reshape(-1, 1)
|
||||
return {
|
||||
"mean": np.mean(values, axis=0).tolist(),
|
||||
"std": np.std(values, axis=0).tolist(),
|
||||
"q01": np.quantile(values, 0.01, axis=0).tolist(),
|
||||
"q99": np.quantile(values, 0.99, axis=0).tolist(),
|
||||
}
|
||||
|
||||
|
||||
def _apply_slice_stats(
|
||||
full_stats: dict[str, list[float]],
|
||||
index_range: list[int],
|
||||
slice_stats: dict[str, list[float]],
|
||||
) -> None:
|
||||
start, end = index_range
|
||||
for field in ("mean", "std", "q01", "q99"):
|
||||
full_stats[field][start:end] = slice_stats[field]
|
||||
|
||||
|
||||
def config_to_index_ranges(config: dict[str, int]) -> dict[str, list[int]]:
|
||||
ranges: dict[str, list[int]] = {}
|
||||
cur = 0
|
||||
for key, dim in config.items():
|
||||
if key in SKIP_DOF_KEYS:
|
||||
continue
|
||||
ranges[key] = [cur, cur + int(dim)]
|
||||
cur += int(dim)
|
||||
return ranges
|
||||
|
||||
|
||||
def load_train_config(path: Path) -> dict[str, Any]:
|
||||
with path.open(encoding="utf-8") as f:
|
||||
config = yaml.load(f, Loader=yaml.SafeLoader)
|
||||
if not isinstance(config, dict):
|
||||
raise ValueError(f"train config must be a YAML mapping, got {type(config)}")
|
||||
return config
|
||||
|
||||
|
||||
def parse_train_config(config: dict[str, Any]) -> TrainConfigContext:
|
||||
task = config.get("task")
|
||||
if not isinstance(task, dict):
|
||||
raise ValueError("train config must contain a 'task' section")
|
||||
|
||||
dof_config = task.get("dof_config")
|
||||
agent_pos_config = task.get("agent_pos_config")
|
||||
if not isinstance(dof_config, dict) or not dof_config:
|
||||
raise ValueError("task.dof_config is required in train config")
|
||||
if not isinstance(agent_pos_config, dict) or not agent_pos_config:
|
||||
raise ValueError("task.agent_pos_config is required in train config")
|
||||
|
||||
data_cfg = config.get("data")
|
||||
if not isinstance(data_cfg, dict):
|
||||
raise ValueError("train config must contain a 'data' section")
|
||||
|
||||
lerobot_config = data_cfg.get("lerobot_config")
|
||||
if not isinstance(lerobot_config, dict):
|
||||
raise ValueError("data.lerobot_config is required in train config")
|
||||
|
||||
repo_id = lerobot_config.get("repo_id")
|
||||
if not repo_id:
|
||||
raise ValueError("data.lerobot_config.repo_id is required in train config")
|
||||
|
||||
norm_stats_path = config.get("norm_stats_path") or data_cfg.get("norm_stats_path")
|
||||
if not norm_stats_path:
|
||||
raise ValueError(
|
||||
"norm_stats_path is required in train config " "(top-level or under data)"
|
||||
)
|
||||
|
||||
key_mappings = data_cfg.get("key_mappings")
|
||||
if not isinstance(key_mappings, dict):
|
||||
raise ValueError("data.key_mappings is required in train config")
|
||||
|
||||
state_key = key_mappings.get("state", "observation.state")
|
||||
action_key = key_mappings.get("action", "action")
|
||||
|
||||
action_chunk = int(
|
||||
task.get("action_horizon")
|
||||
or task.get("action_horizon_flow")
|
||||
or data_cfg.get("action_horizon")
|
||||
or 32
|
||||
)
|
||||
|
||||
return TrainConfigContext(
|
||||
data_root=Path(repo_id),
|
||||
output_path=Path(norm_stats_path),
|
||||
state_key=state_key,
|
||||
action_key=action_key,
|
||||
propri_ranges=config_to_index_ranges(agent_pos_config),
|
||||
action_ranges=config_to_index_ranges(dof_config),
|
||||
action_chunk=action_chunk,
|
||||
dof_config=dof_config,
|
||||
agent_pos_config=agent_pos_config,
|
||||
)
|
||||
|
||||
|
||||
def layout_vector_dim(layout_config: dict[str, int]) -> int:
|
||||
"""Config vector width excluding virtual padding keys."""
|
||||
return sum(
|
||||
int(dim) for key, dim in layout_config.items() if key not in SKIP_DOF_KEYS
|
||||
)
|
||||
|
||||
|
||||
def _prepare_arrays_for_layout(
|
||||
states: np.ndarray,
|
||||
actions: np.ndarray,
|
||||
agent_pos_config: dict[str, int],
|
||||
dof_config: dict[str, int],
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""
|
||||
Compute statistics (min, q01, q99, max) for each action type and dimension.
|
||||
|
||||
Args:
|
||||
action_data_by_robot: Dict[robot_id][action_type] -> list of arrays/lists
|
||||
|
||||
Returns:
|
||||
Dict[robot_id][action_type] -> {
|
||||
"min": [min for each dim],
|
||||
"q01": [quantile 1% for each dim],
|
||||
"q99": [quantile 99% for each dim],
|
||||
"max": [max for each dim],
|
||||
"delta": [max - min for each dim]
|
||||
"delta_q99_q01": [q99 - q01 for each dim]
|
||||
}
|
||||
Match LeRobot training loader: convert 3D Euler slices to 6D when config
|
||||
uses rotation_6D keys but the dataset stores 14-dim Euler vectors.
|
||||
"""
|
||||
stats = {}
|
||||
from wall_x.data.backends.lerobot.rotation_layout import (
|
||||
euler_layout_dim,
|
||||
layout_uses_6d_rotation,
|
||||
maybe_convert_euler_to_6d,
|
||||
)
|
||||
|
||||
for robot_id, action_data in action_data_by_robot.items():
|
||||
stats[robot_id] = {}
|
||||
convert_state = layout_uses_6d_rotation(agent_pos_config)
|
||||
convert_action = layout_uses_6d_rotation(dof_config)
|
||||
|
||||
for action_type, values_list in action_data.items():
|
||||
if not values_list:
|
||||
continue
|
||||
if convert_state:
|
||||
raw_dim = euler_layout_dim(agent_pos_config)
|
||||
target_dim = layout_vector_dim(agent_pos_config)
|
||||
if states.shape[-1] == raw_dim:
|
||||
states = maybe_convert_euler_to_6d(states, agent_pos_config, True)
|
||||
logging.info(
|
||||
"Converted state Euler->6D (%d -> %d dims)", raw_dim, states.shape[-1]
|
||||
)
|
||||
elif states.shape[-1] != target_dim:
|
||||
raise ValueError(
|
||||
f"State dim {states.shape[-1]} does not match Euler raw dim "
|
||||
f"{raw_dim} or 6D layout dim {target_dim} from agent_pos_config"
|
||||
)
|
||||
|
||||
# Convert to numpy array: shape (num_samples, num_dims)
|
||||
try:
|
||||
values_array = np.array(values_list)
|
||||
if values_array.size == 0:
|
||||
continue
|
||||
if convert_action:
|
||||
raw_dim = euler_layout_dim(dof_config)
|
||||
target_dim = layout_vector_dim(dof_config)
|
||||
if actions.shape[-1] == raw_dim:
|
||||
actions = maybe_convert_euler_to_6d(actions, dof_config, True)
|
||||
logging.info(
|
||||
"Converted action Euler->6D (%d -> %d dims)", raw_dim, actions.shape[-1]
|
||||
)
|
||||
elif actions.shape[-1] != target_dim:
|
||||
raise ValueError(
|
||||
f"Action dim {actions.shape[-1]} does not match Euler raw dim "
|
||||
f"{raw_dim} or 6D layout dim {target_dim} from dof_config"
|
||||
)
|
||||
|
||||
# Handle both 1D and 2D cases
|
||||
if values_array.ndim == 1:
|
||||
values_array = values_array.reshape(-1, 1)
|
||||
elif values_array.ndim == 2:
|
||||
pass
|
||||
else:
|
||||
logging.warning(
|
||||
f"Unexpected shape for {robot_id}/{action_type}: {values_array.shape}"
|
||||
)
|
||||
continue
|
||||
expected_state_dim = layout_vector_dim(agent_pos_config)
|
||||
expected_action_dim = layout_vector_dim(dof_config)
|
||||
if states.shape[-1] != expected_state_dim:
|
||||
raise ValueError(
|
||||
f"State dim {states.shape[-1]} != expected layout dim {expected_state_dim}"
|
||||
)
|
||||
if actions.shape[-1] != expected_action_dim:
|
||||
raise ValueError(
|
||||
f"Action dim {actions.shape[-1]} != expected layout dim {expected_action_dim}"
|
||||
)
|
||||
return states, actions
|
||||
|
||||
# Compute statistics for each dimension
|
||||
min_vals = np.min(values_array, axis=0).tolist()
|
||||
max_vals = np.max(values_array, axis=0).tolist()
|
||||
q01_vals = np.quantile(values_array, 0.01, axis=0).tolist()
|
||||
q99_vals = np.quantile(values_array, 0.99, axis=0).tolist()
|
||||
delta_vals = (np.array(max_vals) - np.array(min_vals)).tolist()
|
||||
delta_q99_q01_vals = (np.array(q99_vals) - np.array(q01_vals)).tolist()
|
||||
|
||||
stats[robot_id][action_type] = {
|
||||
"min": min_vals,
|
||||
"q01": q01_vals,
|
||||
"q99": q99_vals,
|
||||
"max": max_vals,
|
||||
"delta": delta_vals,
|
||||
"delta_q99_q01": delta_q99_q01_vals,
|
||||
}
|
||||
def resolve_lerobot_dataset_paths(dataset_root: Path) -> tuple[str, Path]:
|
||||
"""Return ``(repo_id, root)`` for a local LeRobot dataset directory."""
|
||||
root = dataset_root.expanduser().resolve()
|
||||
if not root.is_dir():
|
||||
raise FileNotFoundError(f"LeRobot dataset root not found: {root}")
|
||||
return root.name, root
|
||||
|
||||
except Exception as e:
|
||||
logging.warning(
|
||||
f"Error computing statistics for {robot_id}/{action_type}: {e}"
|
||||
|
||||
def _load_parquet_state_action_only(
|
||||
data_root: Path,
|
||||
state_key: str,
|
||||
action_key: str,
|
||||
):
|
||||
"""
|
||||
Read only state/action columns from LeRobot v3 parquet.
|
||||
|
||||
Does not load image/video columns or decode mp4 files.
|
||||
"""
|
||||
import pyarrow.dataset as pa_ds
|
||||
from datasets import Dataset
|
||||
|
||||
root = data_root.expanduser().resolve()
|
||||
paths = sorted((root / "data").glob("*/*.parquet"))
|
||||
if not paths:
|
||||
raise FileNotFoundError(f"No parquet files under {root / 'data'}")
|
||||
|
||||
logging.info(
|
||||
"Reading parquet columns %r, %r only (no video/images)",
|
||||
state_key,
|
||||
action_key,
|
||||
)
|
||||
arrow_dataset = pa_ds.dataset([str(path) for path in paths], format="parquet")
|
||||
table = arrow_dataset.to_table(columns=[state_key, action_key])
|
||||
return Dataset(table)
|
||||
|
||||
|
||||
def _load_state_action_table(
|
||||
data_root: Path,
|
||||
state_key: str,
|
||||
action_key: str,
|
||||
):
|
||||
root = data_root.expanduser().resolve()
|
||||
if root.is_dir() and (root / "meta" / "info.json").is_file():
|
||||
table = _load_parquet_state_action_only(root, state_key, action_key)
|
||||
return table, state_key, action_key
|
||||
|
||||
try:
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
except ImportError as exc:
|
||||
raise RuntimeError(
|
||||
"compute_norm_stats.py requires LeRobot. Install it first, for example "
|
||||
"`pip install lerobot==0.4.4` or follow the repository README."
|
||||
) from exc
|
||||
|
||||
dataset = LeRobotDataset(str(data_root), root=None, video_backend="pyav")
|
||||
non_image_columns = [
|
||||
col for col in dataset.features if "image" not in col and col not in {"task"}
|
||||
]
|
||||
if state_key not in non_image_columns or action_key not in non_image_columns:
|
||||
raise ValueError(
|
||||
f"Expected keys {state_key!r} and {action_key!r} in dataset columns, "
|
||||
f"got {non_image_columns}"
|
||||
)
|
||||
table = dataset.hf_dataset.select_columns([state_key, action_key])
|
||||
return table, state_key, action_key
|
||||
|
||||
|
||||
def _table_to_arrays(
|
||||
table,
|
||||
state_key: str,
|
||||
action_key: str,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Load full columns once; avoids O(N*chunk) random row access."""
|
||||
logging.info("Loading state/action columns into memory...")
|
||||
try:
|
||||
states = np.asarray(table[state_key], dtype=np.float32)
|
||||
actions = np.asarray(table[action_key], dtype=np.float32)
|
||||
except (KeyError, TypeError, ValueError) as exc:
|
||||
logging.warning(
|
||||
"Column-wise load failed (%s); falling back to per-row stack.", exc
|
||||
)
|
||||
states = np.stack(
|
||||
[
|
||||
np.asarray(table[i][state_key], dtype=np.float32)
|
||||
for i in range(len(table))
|
||||
]
|
||||
)
|
||||
actions = np.stack(
|
||||
[
|
||||
np.asarray(table[i][action_key], dtype=np.float32)
|
||||
for i in range(len(table))
|
||||
]
|
||||
)
|
||||
if states.ndim == 1:
|
||||
states = states.reshape(-1, 1)
|
||||
if actions.ndim == 1:
|
||||
actions = actions.reshape(-1, 1)
|
||||
logging.info(
|
||||
" frames=%d state_dim=%d action_dim=%d",
|
||||
len(states),
|
||||
states.shape[1],
|
||||
actions.shape[1],
|
||||
)
|
||||
return states, actions
|
||||
|
||||
|
||||
def _collect_relative_cartesian(
|
||||
actions: np.ndarray,
|
||||
states: np.ndarray,
|
||||
index_range: list[int],
|
||||
action_chunk: int,
|
||||
) -> np.ndarray:
|
||||
start, end = index_range
|
||||
max_start = max(0, len(actions) - action_chunk)
|
||||
chunks = []
|
||||
anchor_states = states[:max_start, start:end]
|
||||
for offset in range(action_chunk):
|
||||
chunks.append(actions[offset : offset + max_start, start:end] - anchor_states)
|
||||
return np.concatenate(chunks, axis=0)
|
||||
|
||||
|
||||
def _compute_delta_from_state_and_abs_rot(
|
||||
rotations: np.ndarray, state: np.ndarray
|
||||
) -> np.ndarray:
|
||||
"""Relative rotation: R_rel = R_abs @ R_state^T (same convention as the loader)."""
|
||||
_ensure_geometry_funcs()
|
||||
|
||||
if rotations.shape[-1] == 3:
|
||||
rotations_matrix = euler_to_matrix_zyx_batch_nb(rotations)
|
||||
out_is_euler = True
|
||||
elif rotations.shape[-1] == 6:
|
||||
rotations_matrix = so3_to_matrix_batch_nb(rotations)
|
||||
out_is_euler = False
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Only 3D euler or 6D rotation supported, got {rotations.shape[-1]}D"
|
||||
)
|
||||
|
||||
if state.shape[-1] == 3:
|
||||
state_matrix = euler_to_matrix_zyx_batch_nb(state[np.newaxis, :])[0]
|
||||
elif state.shape[-1] == 6:
|
||||
state_matrix = so3_to_matrix_batch_nb(state[np.newaxis, :])[0]
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Only 3D euler or 6D rotation supported, got {state.shape[-1]}D"
|
||||
)
|
||||
|
||||
return _abs_rot_to_delta(rotations_matrix, state_matrix, out_is_euler)
|
||||
|
||||
|
||||
@jit(nopython=True, parallel=True)
|
||||
def _abs_rot_to_delta(
|
||||
rotations_matrix: np.ndarray,
|
||||
state_matrix: np.ndarray,
|
||||
out_is_euler: bool,
|
||||
) -> np.ndarray:
|
||||
st = np.empty((3, 3), dtype=np.float64)
|
||||
st[0, 0] = state_matrix[0, 0]
|
||||
st[0, 1] = state_matrix[1, 0]
|
||||
st[0, 2] = state_matrix[2, 0]
|
||||
st[1, 0] = state_matrix[0, 1]
|
||||
st[1, 1] = state_matrix[1, 1]
|
||||
st[1, 2] = state_matrix[2, 1]
|
||||
st[2, 0] = state_matrix[0, 2]
|
||||
st[2, 1] = state_matrix[1, 2]
|
||||
st[2, 2] = state_matrix[2, 2]
|
||||
|
||||
n = rotations_matrix.shape[0]
|
||||
r_rel = np.empty((n, 3, 3), dtype=np.float64)
|
||||
for i in prange(n):
|
||||
a00 = rotations_matrix[i, 0, 0]
|
||||
a01 = rotations_matrix[i, 0, 1]
|
||||
a02 = rotations_matrix[i, 0, 2]
|
||||
a10 = rotations_matrix[i, 1, 0]
|
||||
a11 = rotations_matrix[i, 1, 1]
|
||||
a12 = rotations_matrix[i, 1, 2]
|
||||
a20 = rotations_matrix[i, 2, 0]
|
||||
a21 = rotations_matrix[i, 2, 1]
|
||||
a22 = rotations_matrix[i, 2, 2]
|
||||
|
||||
r_rel[i, 0, 0] = a00 * st[0, 0] + a01 * st[1, 0] + a02 * st[2, 0]
|
||||
r_rel[i, 0, 1] = a00 * st[0, 1] + a01 * st[1, 1] + a02 * st[2, 1]
|
||||
r_rel[i, 0, 2] = a00 * st[0, 2] + a01 * st[1, 2] + a02 * st[2, 2]
|
||||
r_rel[i, 1, 0] = a10 * st[0, 0] + a11 * st[1, 0] + a12 * st[2, 0]
|
||||
r_rel[i, 1, 1] = a10 * st[0, 1] + a11 * st[1, 1] + a12 * st[2, 1]
|
||||
r_rel[i, 1, 2] = a10 * st[0, 2] + a11 * st[1, 2] + a12 * st[2, 2]
|
||||
r_rel[i, 2, 0] = a20 * st[0, 0] + a21 * st[1, 0] + a22 * st[2, 0]
|
||||
r_rel[i, 2, 1] = a20 * st[0, 1] + a21 * st[1, 1] + a22 * st[2, 1]
|
||||
r_rel[i, 2, 2] = a20 * st[0, 2] + a21 * st[1, 2] + a22 * st[2, 2]
|
||||
|
||||
if out_is_euler:
|
||||
d_euler = matrix_to_euler_zyx_batch_nb(r_rel)
|
||||
return canonicalize_euler_zyx_batch_nb(d_euler)
|
||||
|
||||
out6 = np.empty((n, 6), dtype=np.float64)
|
||||
for i in prange(n):
|
||||
out6[i, 0] = r_rel[i, 0, 0]
|
||||
out6[i, 1] = r_rel[i, 0, 1]
|
||||
out6[i, 2] = r_rel[i, 0, 2]
|
||||
out6[i, 3] = r_rel[i, 1, 0]
|
||||
out6[i, 4] = r_rel[i, 1, 1]
|
||||
out6[i, 5] = r_rel[i, 1, 2]
|
||||
return out6
|
||||
|
||||
|
||||
def _collect_relative_rotation(
|
||||
actions: np.ndarray,
|
||||
states: np.ndarray,
|
||||
index_range: list[int],
|
||||
action_chunk: int,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Per-anchor action chunk relative to anchor state (matches lerobot loader).
|
||||
|
||||
Unlike cartesian relative, each anchor processes a full [chunk, dim] action
|
||||
clip against a single proprio rotation at the anchor frame.
|
||||
"""
|
||||
start, end = index_range
|
||||
max_start = max(0, len(actions) - action_chunk)
|
||||
if max_start == 0:
|
||||
return np.empty((0, end - start), dtype=np.float32)
|
||||
|
||||
chunks = []
|
||||
for anchor_idx in tqdm(
|
||||
range(max_start),
|
||||
desc=" relative rotation anchors",
|
||||
leave=False,
|
||||
):
|
||||
action_clip = actions[anchor_idx : anchor_idx + action_chunk, start:end]
|
||||
proprio_clip = states[anchor_idx, start:end]
|
||||
rel = _compute_delta_from_state_and_abs_rot(
|
||||
action_clip.astype(np.float64), proprio_clip.astype(np.float64)
|
||||
).astype(np.float32)
|
||||
chunks.append(rel)
|
||||
return np.concatenate(chunks, axis=0)
|
||||
|
||||
|
||||
def collect_dof_vectors_from_arrays(
|
||||
states: np.ndarray,
|
||||
actions: np.ndarray,
|
||||
propri_ranges: dict[str, list[int]],
|
||||
action_ranges: dict[str, list[int]],
|
||||
action_chunk: int = 32,
|
||||
) -> dict[str, np.ndarray]:
|
||||
vectors: dict[str, np.ndarray] = {}
|
||||
|
||||
absolute_action_keys = {
|
||||
key: index_range
|
||||
for key, index_range in action_ranges.items()
|
||||
if not key.endswith("_relative")
|
||||
}
|
||||
relative_action_keys = {
|
||||
key: index_range
|
||||
for key, index_range in action_ranges.items()
|
||||
if key.endswith("_relative")
|
||||
}
|
||||
|
||||
for sub_key, index_range in propri_ranges.items():
|
||||
start, end = index_range
|
||||
vectors[sub_key] = states[:, start:end]
|
||||
|
||||
for sub_key, index_range in absolute_action_keys.items():
|
||||
start, end = index_range
|
||||
vectors[sub_key] = actions[:, start:end]
|
||||
|
||||
if relative_action_keys:
|
||||
logging.info(
|
||||
"Computing relative action slices (chunk=%d, anchors=%d)...",
|
||||
action_chunk,
|
||||
max(0, len(actions) - action_chunk),
|
||||
)
|
||||
for sub_key, index_range in tqdm(
|
||||
relative_action_keys.items(), desc="Relative action keys"
|
||||
):
|
||||
if "rotation" in sub_key:
|
||||
vectors[sub_key] = _collect_relative_rotation(
|
||||
actions, states, index_range, action_chunk
|
||||
)
|
||||
else:
|
||||
vectors[sub_key] = _collect_relative_cartesian(
|
||||
actions, states, index_range, action_chunk
|
||||
)
|
||||
continue
|
||||
|
||||
return stats
|
||||
return vectors
|
||||
|
||||
|
||||
def load_lerobot_dataset(
|
||||
repo_id: str,
|
||||
trajectory_keys: Dict,
|
||||
base_dir: Path,
|
||||
) -> None:
|
||||
def compute_norm_stats_with_dof_config(
|
||||
data_root: Path,
|
||||
output_path: Path,
|
||||
propri_ranges: dict[str, list[int]],
|
||||
action_ranges: dict[str, list[int]],
|
||||
state_key: str = "observation.state",
|
||||
action_key: str = "action",
|
||||
action_chunk: int = 32,
|
||||
dof_config: dict[str, int] | None = None,
|
||||
agent_pos_config: dict[str, int] | None = None,
|
||||
) -> dict[str, dict]:
|
||||
table, state_key, action_key = _load_state_action_table(
|
||||
data_root, state_key, action_key
|
||||
)
|
||||
states, actions = _table_to_arrays(table, state_key, action_key)
|
||||
if dof_config is not None and agent_pos_config is not None:
|
||||
states, actions = _prepare_arrays_for_layout(
|
||||
states, actions, agent_pos_config, dof_config
|
||||
)
|
||||
norm_stats = {
|
||||
state_key: compute_vector_stats(states),
|
||||
action_key: compute_vector_stats(actions),
|
||||
}
|
||||
|
||||
# Load local or remote dataset
|
||||
dataset = LeRobotDataset(base_dir)
|
||||
vectors = collect_dof_vectors_from_arrays(
|
||||
states=states,
|
||||
actions=actions,
|
||||
propri_ranges=propri_ranges,
|
||||
action_ranges=action_ranges,
|
||||
action_chunk=action_chunk,
|
||||
)
|
||||
|
||||
# Iterate through all data
|
||||
frames: Dict[str, Dict[str, List]] = defaultdict(lambda: defaultdict(list))
|
||||
for sub_key, index_range in propri_ranges.items():
|
||||
if sub_key not in vectors:
|
||||
logging.warning("No samples collected for propri key %s, skipping", sub_key)
|
||||
continue
|
||||
slice_stats = compute_vector_stats(vectors[sub_key])
|
||||
_apply_slice_stats(norm_stats[state_key], index_range, slice_stats)
|
||||
logging.info(" %s (agent_pos): dim=%d", sub_key, len(slice_stats["mean"]))
|
||||
|
||||
all_features = dataset.features
|
||||
non_image_columns = [col for col in all_features if "image" not in col]
|
||||
for sub_key, index_range in action_ranges.items():
|
||||
if sub_key not in vectors:
|
||||
logging.warning("No samples collected for action key %s, skipping", sub_key)
|
||||
continue
|
||||
slice_stats = compute_vector_stats(vectors[sub_key])
|
||||
_apply_slice_stats(norm_stats[action_key], index_range, slice_stats)
|
||||
mode = "relative" if sub_key.endswith("_relative") else "absolute"
|
||||
logging.info(" %s (dof, %s): dim=%d", sub_key, mode, len(slice_stats["mean"]))
|
||||
|
||||
print(f"Reading the following fields:{non_image_columns}")
|
||||
fast_dataset = dataset.hf_dataset.select_columns(non_image_columns)
|
||||
|
||||
for i in tqdm(range(len(fast_dataset))):
|
||||
sample = fast_dataset[i]
|
||||
action = sample["action"] # torch.Tensor
|
||||
propri = sample["observation.state"]
|
||||
|
||||
for key, action_keys in trajectory_keys.items():
|
||||
for action_key, action_range in action_keys.items():
|
||||
if key == "action":
|
||||
frames[repo_id][action_key].append(
|
||||
action[action_range[0] : action_range[1]].numpy().tolist()
|
||||
)
|
||||
else:
|
||||
frames[repo_id][action_key].append(
|
||||
propri[action_range[0] : action_range[1]].numpy().tolist()
|
||||
)
|
||||
|
||||
return frames
|
||||
write_norm_stats(output_path, norm_stats)
|
||||
return norm_stats
|
||||
|
||||
|
||||
def compute_action_normalizer(
|
||||
repo_id: str, trajectory_keys: Dict, base_dir: Path, output_dir: Path
|
||||
) -> None:
|
||||
"""
|
||||
Compute action normalizer statistics for all robot_ids.
|
||||
"""
|
||||
logging.info("Starting action normalizer computation...")
|
||||
def load_vectors(
|
||||
data_root: Path,
|
||||
state_key: str,
|
||||
action_key: str,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
table, state_key, action_key = _load_state_action_table(
|
||||
data_root, state_key, action_key
|
||||
)
|
||||
return _table_to_arrays(table, state_key, action_key)
|
||||
|
||||
frames = load_lerobot_dataset(repo_id, trajectory_keys, base_dir)
|
||||
|
||||
# Compute statistics
|
||||
stats = compute_action_statistics(frames)
|
||||
def compute_norm_stats(
|
||||
data_root: Path,
|
||||
output_path: Path,
|
||||
state_key: str = "observation.state",
|
||||
action_key: str = "action",
|
||||
train_ctx: TrainConfigContext | None = None,
|
||||
) -> dict[str, dict]:
|
||||
if train_ctx is None:
|
||||
states, actions = load_vectors(data_root, state_key, action_key)
|
||||
norm_stats = {
|
||||
state_key: compute_vector_stats(states),
|
||||
action_key: compute_vector_stats(actions),
|
||||
}
|
||||
write_norm_stats(output_path, norm_stats)
|
||||
return norm_stats
|
||||
|
||||
# Save statistics for each robot_id
|
||||
output_dir = Path(output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
return compute_norm_stats_with_dof_config(
|
||||
data_root=data_root,
|
||||
output_path=output_path,
|
||||
propri_ranges=train_ctx.propri_ranges,
|
||||
action_ranges=train_ctx.action_ranges,
|
||||
state_key=state_key,
|
||||
action_key=action_key,
|
||||
action_chunk=train_ctx.action_chunk,
|
||||
dof_config=train_ctx.dof_config,
|
||||
agent_pos_config=train_ctx.agent_pos_config,
|
||||
)
|
||||
|
||||
# for robot_id, robot_stats in stats.items():
|
||||
# output_file = output_dir / f"{robot_id}_action_stats.json"
|
||||
# write_json(output_file, robot_stats)
|
||||
# logging.info(f"Saved action statistics for {robot_id} to {output_file}")
|
||||
|
||||
# Also save a combined file
|
||||
combined_output = output_dir / "all_robots_action_stats.json"
|
||||
write_json(combined_output, stats)
|
||||
logging.info(f"Saved combined action statistics to {combined_output}")
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Compute norm stats for a local LeRobot v3 dataset.",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""\
|
||||
examples:
|
||||
%(prog)s --train_config /path/to/train_config.yml
|
||||
%(prog)s --train_config /path/to/multitask_config.yml \\
|
||||
--data_root /path/to/repo_id --output_path /path/to/norm_stats_path
|
||||
%(prog)s --data_root /path/to/lerobot_dataset --output_path /path/to/norm_stats.json
|
||||
""",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--train_config",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Training YAML config (e.g. cvpr_example.yml). When set, reads "
|
||||
"data.lerobot_config.repo_id, data.norm_stats_path, task.dof_config, "
|
||||
"task.agent_pos_config and task.action_horizon. Action keys ending "
|
||||
"with '_relative' use the same relative-pose logic as lerobot loader."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--data_root",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Local LeRobot dataset directory (overrides train config)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Output json path (overrides train config)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--state_key",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Dataset column for proprioception (overrides train config)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--action_key",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Dataset column for action (overrides train config)",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
|
||||
args = parse_args()
|
||||
|
||||
repo_id = "xxx" # your dataset name
|
||||
data_root_path = "/path/to/lerobot/dataset"
|
||||
output_stats_dir = "/path/to/save/action_stats"
|
||||
trajectory_keys = { # your dataset keys
|
||||
"propri": {
|
||||
"follow_right_ee_cartesian_pos": [0, 3],
|
||||
"follow_right_ee_rotation": [3, 6],
|
||||
"follow_right_gripper": [6, 7],
|
||||
},
|
||||
"action": {
|
||||
"master_right_ee_cartesian_pos": [0, 3],
|
||||
"master_right_ee_rotation": [3, 6],
|
||||
"master_right_gripper": [6, 7],
|
||||
},
|
||||
}
|
||||
train_ctx: TrainConfigContext | None = None
|
||||
if args.train_config:
|
||||
config_path = Path(args.train_config)
|
||||
if not config_path.exists():
|
||||
raise FileNotFoundError(f"train config not found: {config_path}")
|
||||
train_ctx = parse_train_config(load_train_config(config_path))
|
||||
logging.info("train_config: %s", config_path)
|
||||
logging.info(" dof_config keys: %s", list(train_ctx.action_ranges))
|
||||
logging.info(" agent_pos_config keys: %s", list(train_ctx.propri_ranges))
|
||||
logging.info(" action_chunk: %d", train_ctx.action_chunk)
|
||||
|
||||
compute_action_normalizer(
|
||||
repo_id, trajectory_keys, data_root_path, output_stats_dir
|
||||
data_root = Path(args.data_root or (train_ctx.data_root if train_ctx else ""))
|
||||
|
||||
output_path = Path(args.output_path or (train_ctx.output_path if train_ctx else ""))
|
||||
state_key = args.state_key or (
|
||||
train_ctx.state_key if train_ctx else "observation.state"
|
||||
)
|
||||
logging.info("Action normalizer computation completed.")
|
||||
action_key = args.action_key or (train_ctx.action_key if train_ctx else "action")
|
||||
|
||||
if not data_root.exists():
|
||||
raise FileNotFoundError(f"Dataset not found: {data_root}")
|
||||
|
||||
logging.info("dataset: %s", data_root)
|
||||
logging.info("output: %s", output_path)
|
||||
|
||||
norm_stats = compute_norm_stats(
|
||||
data_root=data_root,
|
||||
output_path=output_path,
|
||||
state_key=state_key,
|
||||
action_key=action_key,
|
||||
train_ctx=train_ctx,
|
||||
)
|
||||
|
||||
for key, stats in norm_stats.items():
|
||||
logging.info(" %s: dim=%d", key, len(stats["mean"]))
|
||||
|
||||
logging.info("Saved norm stats to %s", output_path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
+1098
-117
File diff suppressed because it is too large
Load Diff
Regular → Executable
+109
-70
@@ -1,85 +1,124 @@
|
||||
import torch
|
||||
from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction
|
||||
#!/usr/bin/env python3
|
||||
"""Run one Wall-X VLA inference pass through the harrix adapter.
|
||||
|
||||
model_path = "/path/to/model"
|
||||
model = Qwen2_5_VLMoEForAction.from_pretrained(model_path)
|
||||
model.eval()
|
||||
This is a lightweight smoke test for the inference path. It builds a synthetic
|
||||
LIBERO-style observation, loads the checkpoint through harrix, and prints the
|
||||
predicted action chunk shape.
|
||||
"""
|
||||
|
||||
# Gen Fake data
|
||||
batch_size = 1
|
||||
seq_length = 50
|
||||
from __future__ import annotations
|
||||
|
||||
torch.manual_seed(0)
|
||||
fake_input_ids = torch.randint(
|
||||
0, len(model.processor.tokenizer), (batch_size, seq_length), dtype=torch.long
|
||||
)
|
||||
fake_attention_mask = torch.ones((batch_size, seq_length), dtype=torch.long)
|
||||
fake_moe_token_types = torch.zeros((batch_size, seq_length), dtype=torch.long)
|
||||
fake_position_ids = (
|
||||
torch.arange(seq_length, dtype=torch.long).unsqueeze(0).expand(batch_size, -1)
|
||||
)
|
||||
fake_proprioception = torch.randn((batch_size, 1, 20), dtype=torch.float32)
|
||||
fake_agent_pos_mask = torch.ones((batch_size, 1, 20), dtype=torch.float32)
|
||||
fake_dof_mask = torch.ones((batch_size, 32, 20), dtype=torch.float32)
|
||||
fake_dataset_names = ["x2_normal"]
|
||||
import argparse
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
device = "cuda"
|
||||
def _ensure_local_harrix_on_path() -> None:
|
||||
repo_root = Path(__file__).resolve().parents[1]
|
||||
harrix_python = repo_root / "third_party" / "harrix" / "python"
|
||||
if harrix_python.is_dir():
|
||||
sys.path.insert(0, str(harrix_python))
|
||||
|
||||
model = model.to(device)
|
||||
model = model.bfloat16()
|
||||
|
||||
fake_input_ids = fake_input_ids.to(device)
|
||||
fake_attention_mask = fake_attention_mask.to(device)
|
||||
fake_moe_token_types = fake_moe_token_types.to(device)
|
||||
fake_position_ids = fake_position_ids.to(device)
|
||||
fake_proprioception = fake_proprioception.to(device).bfloat16()
|
||||
fake_agent_pos_mask = fake_agent_pos_mask.to(device).bfloat16()
|
||||
fake_dof_mask = fake_dof_mask.to(device).bfloat16()
|
||||
def _right_gripper_dim_from_config(train_config: dict) -> int:
|
||||
layout = train_config.get("agent_pos_config") or train_config.get("task", {}).get(
|
||||
"agent_pos_config", {}
|
||||
)
|
||||
if not isinstance(layout, dict):
|
||||
return 1
|
||||
|
||||
try:
|
||||
with torch.no_grad():
|
||||
outputs = model(
|
||||
input_ids=fake_input_ids,
|
||||
attention_mask=fake_attention_mask,
|
||||
moe_token_types=fake_moe_token_types,
|
||||
position_ids=fake_position_ids,
|
||||
proprioception=fake_proprioception,
|
||||
agent_pos_mask=fake_agent_pos_mask,
|
||||
dof_mask=fake_dof_mask,
|
||||
dataset_names=fake_dataset_names,
|
||||
mode="validate",
|
||||
)
|
||||
gripper_dim = 1
|
||||
for key, dim in layout.items():
|
||||
bare = key.replace("follow_", "").replace("master_", "")
|
||||
if bare == "right_gripper":
|
||||
gripper_dim = int(dim)
|
||||
break
|
||||
|
||||
print("✅ Fake inference test successful!")
|
||||
print(f"Output logits shape: {outputs.logits.shape}")
|
||||
print(f"Output logits dtype: {outputs.logits.dtype}")
|
||||
print(f"Output logits device: {outputs.logits.device}")
|
||||
norm_dim = int(train_config.get("_libero_proprio_norm_dim") or 0)
|
||||
real_dim = sum(int(v) for k, v in layout.items() if k != "action_padding")
|
||||
if norm_dim == real_dim + 1:
|
||||
gripper_dim += 1
|
||||
return max(1, gripper_dim)
|
||||
|
||||
# Check if output is reasonable
|
||||
if outputs.logits.shape == (batch_size, seq_length, model.config.vocab_size):
|
||||
print("✅ Output shape correct")
|
||||
else:
|
||||
print("❌ Output shape incorrect")
|
||||
|
||||
if not torch.isnan(outputs.logits).any():
|
||||
print("✅ Output contains no NaN values")
|
||||
else:
|
||||
print("❌ Output contains NaN values")
|
||||
def _build_fake_observation(seed: int, image_size: int, gripper_dim: int) -> dict:
|
||||
rng = np.random.default_rng(seed)
|
||||
return {
|
||||
"eef_pos": rng.normal(size=(3,)).astype(np.float32),
|
||||
"eef_axisangle": rng.normal(size=(3,)).astype(np.float32),
|
||||
"gripper": rng.normal(size=(gripper_dim,)).astype(np.float32),
|
||||
"face_view": rng.integers(0, 256, (image_size, image_size, 3), dtype=np.uint8),
|
||||
"wrist_view": rng.integers(0, 256, (image_size, image_size, 3), dtype=np.uint8),
|
||||
}
|
||||
|
||||
if not torch.isinf(outputs.logits).any():
|
||||
print("✅ Output contains no infinity values")
|
||||
else:
|
||||
print("❌ Output contains infinity values")
|
||||
|
||||
print("Output logits statistics:")
|
||||
print(f" Min value: {outputs.logits.min().item():.4f}")
|
||||
print(f" Max value: {outputs.logits.max().item():.4f}")
|
||||
print(f" Mean: {outputs.logits.mean().item():.4f}")
|
||||
print(f" Standard deviation: {outputs.logits.std().item():.4f}")
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--checkpoint-path", required=True, help="Checkpoint directory."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--train-config-path",
|
||||
default=None,
|
||||
help="Optional training config path. Defaults to config.yml/config.yaml next to the checkpoint.",
|
||||
)
|
||||
parser.add_argument("--norm-key", default="libero_all")
|
||||
parser.add_argument("--architecture", default="qwen2_5")
|
||||
parser.add_argument("--action-mode", default="flow")
|
||||
parser.add_argument(
|
||||
"--cam-names",
|
||||
nargs="+",
|
||||
default=["face_view", "right_wrist_view"],
|
||||
help="Camera names expected by the checkpoint.",
|
||||
)
|
||||
parser.add_argument("--action-horizon", type=int, default=None)
|
||||
parser.add_argument("--instruction", default="pick up the object")
|
||||
parser.add_argument("--image-size", type=int, default=128)
|
||||
parser.add_argument("--seed", type=int, default=0)
|
||||
return parser.parse_args()
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Fake inference test failed: {e}")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
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 (
|
||||
EvalConfig,
|
||||
LiberoEnvParams,
|
||||
autofill_from_checkpoint,
|
||||
)
|
||||
|
||||
cfg = EvalConfig()
|
||||
cfg.model.checkpoint_path = args.checkpoint_path
|
||||
cfg.model.train_config_path = args.train_config_path
|
||||
cfg.model.norm_key = args.norm_key
|
||||
cfg.model.cam_names = list(args.cam_names)
|
||||
cfg.model.action_horizon = args.action_horizon
|
||||
cfg.model.architecture = args.architecture
|
||||
cfg.model.action_mode = args.action_mode
|
||||
cfg.env.libero = LiberoEnvParams(num_trials_per_task=1, task_indices=[0])
|
||||
cfg = autofill_from_checkpoint(cfg)
|
||||
|
||||
adapter = build_adapter(cfg)
|
||||
gripper_dim = _right_gripper_dim_from_config(getattr(adapter, "_train_config", {}))
|
||||
payload = {
|
||||
"observation": _build_fake_observation(args.seed, args.image_size, gripper_dim),
|
||||
"instruction": args.instruction,
|
||||
"noise": None,
|
||||
}
|
||||
actions = adapter.predict_batch([payload])
|
||||
action = np.asarray(actions[0])
|
||||
|
||||
print("Fake inference succeeded.")
|
||||
print(f"action shape: {action.shape}")
|
||||
print(f"action dtype: {action.dtype}")
|
||||
print(f"action min/max: {float(action.min()):.6f} / {float(action.max()):.6f}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
|
||||
Regular → Executable
+231
-218
@@ -1,222 +1,235 @@
|
||||
import argparse
|
||||
import time
|
||||
import os
|
||||
#!/usr/bin/env python3
|
||||
"""Run LIBERO evaluation through harrix.
|
||||
|
||||
# from wall_x.utils.baseline_utils import check_baseline_dump, update_baseline
|
||||
from wall_x.infer.utils_libero import set_seed_everywhere, TaskSuite, TASK_MAX_STEPS
|
||||
from wall_x.infer.infer_config import InferConfig
|
||||
from wall_x.infer.env_libero import LiberoRobotEnv
|
||||
The script accepts either a full harrix EvalConfig YAML or a checkpoint path
|
||||
plus common command-line overrides. It intentionally bypasses the legacy
|
||||
Wall-X inference stack.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
|
||||
LIBERO_DEFAULT_MAX_INFER_TIMES = {
|
||||
"libero_spatial": 22,
|
||||
"libero_object": 28,
|
||||
"libero_goal": 30,
|
||||
"libero_10": 52,
|
||||
"libero_90": 40,
|
||||
}
|
||||
|
||||
|
||||
def _ensure_local_harrix_on_path() -> None:
|
||||
repo_root = Path(__file__).resolve().parents[1]
|
||||
harrix_python = repo_root / "third_party" / "harrix" / "python"
|
||||
if harrix_python.is_dir():
|
||||
sys.path.insert(0, str(harrix_python))
|
||||
|
||||
|
||||
def _parse_task_indices(value: str | None) -> list[int] | None:
|
||||
if value is None or value.strip() == "":
|
||||
return None
|
||||
return [int(x) for x in value.split(",") if x.strip()]
|
||||
|
||||
|
||||
def _resolve_max_infer_times(
|
||||
task_suite_name: str | None, max_infer_times: int | None
|
||||
) -> int:
|
||||
if max_infer_times is not None:
|
||||
return max_infer_times
|
||||
suite = task_suite_name or "libero_spatial"
|
||||
return LIBERO_DEFAULT_MAX_INFER_TIMES.get(suite, 22)
|
||||
|
||||
|
||||
def _load_or_build_raw_config(args: argparse.Namespace) -> dict:
|
||||
if args.config is not None:
|
||||
with open(args.config, "r") as f:
|
||||
raw = yaml.safe_load(f) or {}
|
||||
model = raw.setdefault("model", {})
|
||||
env = raw.setdefault("env", {})
|
||||
libero = env.setdefault("libero", {})
|
||||
runtime = raw.setdefault("runtime", {})
|
||||
debug = raw.setdefault("debug", {})
|
||||
if args.checkpoint_path is not None:
|
||||
model["checkpoint_path"] = args.checkpoint_path
|
||||
if args.train_config_path is not None:
|
||||
model["train_config_path"] = args.train_config_path
|
||||
task_indices = _parse_task_indices(args.task_indices)
|
||||
if task_indices is not None:
|
||||
libero["task_indices"] = task_indices
|
||||
if args.max_infer_times is not None or libero.get("max_infer_times") is None:
|
||||
libero["max_infer_times"] = _resolve_max_infer_times(
|
||||
libero.get("task_suite_name", args.task_suite_name),
|
||||
args.max_infer_times,
|
||||
)
|
||||
if args.smoke:
|
||||
libero["task_indices"] = [0]
|
||||
libero["num_trials_per_task"] = 5
|
||||
runtime["num_workers"] = 1
|
||||
runtime["max_batch_size"] = 1
|
||||
if args.deterministic_model:
|
||||
debug["deterministic_model"] = True
|
||||
return raw
|
||||
else:
|
||||
if args.checkpoint_path is None:
|
||||
raise ValueError("--checkpoint-path is required when --config is not set")
|
||||
max_infer_times = _resolve_max_infer_times(
|
||||
args.task_suite_name, args.max_infer_times
|
||||
)
|
||||
raw = {
|
||||
"model": {
|
||||
"checkpoint_path": args.checkpoint_path,
|
||||
"norm_key": args.norm_key,
|
||||
"cam_names": args.cam_names,
|
||||
"architecture": args.architecture,
|
||||
"action_mode": args.action_mode,
|
||||
},
|
||||
"env": {
|
||||
"type": "libero",
|
||||
"seed": args.seed,
|
||||
"libero": {
|
||||
"task_suite_name": args.task_suite_name,
|
||||
"initial_states_path": args.initial_states_path,
|
||||
"num_trials_per_task": args.num_trials_per_task,
|
||||
"max_infer_times": max_infer_times,
|
||||
"skip_intermediate_render": args.skip_intermediate_render,
|
||||
},
|
||||
},
|
||||
"runtime": {
|
||||
"num_workers": args.num_workers,
|
||||
"max_batch_size": args.max_batch_size,
|
||||
"ws_port": args.ws_port,
|
||||
"log_dir": args.log_dir,
|
||||
"driver_mode": args.driver_mode,
|
||||
},
|
||||
"debug": {"deterministic_model": args.deterministic_model},
|
||||
}
|
||||
|
||||
model = raw.setdefault("model", {})
|
||||
env = raw.setdefault("env", {})
|
||||
libero = env.setdefault("libero", {})
|
||||
runtime = raw.setdefault("runtime", {})
|
||||
debug = raw.setdefault("debug", {})
|
||||
|
||||
if args.checkpoint_path is not None:
|
||||
model["checkpoint_path"] = args.checkpoint_path
|
||||
if args.train_config_path is not None:
|
||||
model["train_config_path"] = args.train_config_path
|
||||
if args.norm_key is not None:
|
||||
model["norm_key"] = args.norm_key
|
||||
if args.cam_names is not None:
|
||||
model["cam_names"] = args.cam_names
|
||||
if args.action_horizon is not None:
|
||||
model["action_horizon"] = args.action_horizon
|
||||
if args.architecture is not None:
|
||||
model["architecture"] = args.architecture
|
||||
if args.action_mode is not None:
|
||||
model["action_mode"] = args.action_mode
|
||||
|
||||
env["type"] = "libero"
|
||||
env["seed"] = args.seed
|
||||
libero["task_suite_name"] = args.task_suite_name
|
||||
libero["initial_states_path"] = args.initial_states_path
|
||||
libero["num_trials_per_task"] = args.num_trials_per_task
|
||||
libero["max_infer_times"] = _resolve_max_infer_times(
|
||||
args.task_suite_name, args.max_infer_times
|
||||
)
|
||||
libero["skip_intermediate_render"] = args.skip_intermediate_render
|
||||
task_indices = _parse_task_indices(args.task_indices)
|
||||
if task_indices is not None:
|
||||
libero["task_indices"] = task_indices
|
||||
if args.smoke:
|
||||
libero["task_indices"] = [0]
|
||||
libero["num_trials_per_task"] = 5
|
||||
runtime["num_workers"] = 1
|
||||
runtime["max_batch_size"] = 1
|
||||
|
||||
runtime["num_workers"] = (
|
||||
args.num_workers if not args.smoke else runtime["num_workers"]
|
||||
)
|
||||
runtime["max_batch_size"] = (
|
||||
args.max_batch_size if not args.smoke else runtime["max_batch_size"]
|
||||
)
|
||||
runtime["ws_port"] = args.ws_port
|
||||
runtime["log_dir"] = args.log_dir
|
||||
runtime["driver_mode"] = args.driver_mode
|
||||
debug["deterministic_model"] = args.deterministic_model
|
||||
return raw
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--config", default=None, help="Optional harrix EvalConfig YAML."
|
||||
)
|
||||
parser.add_argument("--checkpoint-path", default=None)
|
||||
parser.add_argument("--train-config-path", default=None)
|
||||
parser.add_argument("--norm-key", default="libero_all")
|
||||
parser.add_argument("--architecture", default="qwen2_5")
|
||||
parser.add_argument("--action-mode", default="flow")
|
||||
parser.add_argument(
|
||||
"--cam-names", nargs="+", default=["face_view", "right_wrist_view"]
|
||||
)
|
||||
parser.add_argument("--action-horizon", type=int, default=None)
|
||||
parser.add_argument("--task-suite-name", default="libero_spatial")
|
||||
parser.add_argument("--initial-states-path", default="DEFAULT")
|
||||
parser.add_argument("--num-trials-per-task", type=int, default=50)
|
||||
parser.add_argument(
|
||||
"--task-indices", default=None, help="Comma-separated task ids."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-infer-times",
|
||||
type=int,
|
||||
default=None,
|
||||
help=(
|
||||
"Number of model action chunks per episode. Defaults are suite-specific "
|
||||
"and match the internal LIBERO evaluator."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--num-workers", type=int, default=1)
|
||||
parser.add_argument("--max-batch-size", type=int, default=1)
|
||||
parser.add_argument("--ws-port", type=int, default=8765)
|
||||
parser.add_argument("--log-dir", default="/tmp/harrix_libero_eval")
|
||||
parser.add_argument("--driver-mode", choices=["in_process"], default="in_process")
|
||||
parser.add_argument("--smoke", action="store_true")
|
||||
parser.add_argument("--deterministic-model", action="store_true")
|
||||
parser.add_argument(
|
||||
"--skip-intermediate-render",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
raw = _load_or_build_raw_config(args)
|
||||
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as f:
|
||||
yaml.safe_dump(raw, f, sort_keys=False)
|
||||
tmp_config = f.name
|
||||
|
||||
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
|
||||
|
||||
run(cfg)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = argparse.ArgumentParser(description="Wall-X Libero evaluation script")
|
||||
args.add_argument("--seed", type=int, default=42, help="Random seed")
|
||||
args.add_argument("--id", type=int, default=None, help="Unique index id")
|
||||
args.add_argument("--name", type=int, default=None, help="Launch command name")
|
||||
args.add_argument(
|
||||
"--baseline_path", type=str, default=None, help="Path to baseline record table"
|
||||
)
|
||||
args.add_argument(
|
||||
"--update_baseline",
|
||||
type=bool,
|
||||
default=False,
|
||||
help="Whether to update baseline table",
|
||||
)
|
||||
args.add_argument(
|
||||
"--mode", type=str, default="flow", choices=["flow", "ar"], help="Running mode"
|
||||
)
|
||||
args.add_argument(
|
||||
"--checkpoint_path", type=str, required=True, help="Model checkpoint path"
|
||||
)
|
||||
args.add_argument(
|
||||
"--train_config_path",
|
||||
type=str,
|
||||
required=False,
|
||||
default=None,
|
||||
help="Path to training config .yml file",
|
||||
)
|
||||
args.add_argument(
|
||||
"--norm_key",
|
||||
type=str,
|
||||
default="physical-intelligence/libero",
|
||||
help="Key for normalization statistics",
|
||||
)
|
||||
args.add_argument(
|
||||
"--cam_names",
|
||||
nargs="+",
|
||||
default=["face_view", "right_wrist_view"],
|
||||
help="List of camera names (e.g., --cam_names face_view right_wrist_view)",
|
||||
)
|
||||
args.add_argument(
|
||||
"--task_suite_name",
|
||||
type=str,
|
||||
default=TaskSuite.LIBERO_SPATIAL,
|
||||
choices=[e.value for e in TaskSuite],
|
||||
help="Libero task suite to load",
|
||||
)
|
||||
args.add_argument(
|
||||
"--initial_states_path",
|
||||
type=str,
|
||||
default="DEFAULT",
|
||||
help="Path to initial states .json file, or 'DEFAULT' to use default states.",
|
||||
)
|
||||
args.add_argument(
|
||||
"--num_trials_per_task",
|
||||
type=int,
|
||||
default=50,
|
||||
help="Number of evaluation episodes to run per task",
|
||||
)
|
||||
args.add_argument(
|
||||
"--rollout_dir",
|
||||
type=str,
|
||||
default="./rollouts",
|
||||
help="Directory to save rollout videos",
|
||||
)
|
||||
args = args.parse_args()
|
||||
|
||||
print(f"Using random seed: {args.seed}")
|
||||
set_seed_everywhere(args.seed)
|
||||
|
||||
print("Initializing InferConfig...")
|
||||
if args.train_config_path is None:
|
||||
args.train_config_path = os.path.join(args.checkpoint_path, "config.yml")
|
||||
|
||||
config = InferConfig(
|
||||
checkpoint_path=args.checkpoint_path,
|
||||
train_config_path=args.train_config_path,
|
||||
norm_key=args.norm_key,
|
||||
cam_names=args.cam_names,
|
||||
)
|
||||
if args.mode == "flow":
|
||||
config.action_horizon = config.train_config.get("data", {}).get(
|
||||
"action_horizon_flow", 10
|
||||
)
|
||||
elif args.mode == "ar":
|
||||
config.action_horizon = config.train_config.get("data", {}).get(
|
||||
"action_horizon_ar", 10
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Invalid mode: {args.mode}")
|
||||
config.model_device = "cuda"
|
||||
|
||||
print("Initializing LiberoRobotEnv (Evaluator)...")
|
||||
|
||||
config.action_dim = 7
|
||||
config.pred_horizon = 10
|
||||
|
||||
evaluator = LiberoRobotEnv(
|
||||
config=config,
|
||||
task_suite_name=args.task_suite_name,
|
||||
initial_states_path=args.initial_states_path,
|
||||
rollout_dir=args.rollout_dir,
|
||||
seed=args.seed,
|
||||
)
|
||||
|
||||
print(f"\n{'='*20} Starting Evaluation {'='*20}")
|
||||
print(f"Task suite: {args.task_suite_name}")
|
||||
print(f"Number of tasks: {evaluator.num_tasks}")
|
||||
print(f"Trials per task: {args.num_trials_per_task}")
|
||||
print(f"Initial states: {args.initial_states_path}")
|
||||
print(f"Videos will be saved to: {evaluator.rollout_dir}")
|
||||
print(f"{'='*50}\n")
|
||||
|
||||
total_successes = 0
|
||||
total_episodes_run = 0
|
||||
start_time = time.time()
|
||||
|
||||
for task_id in range(evaluator.num_tasks):
|
||||
task_successes = 0
|
||||
task_episodes_attempted = 0
|
||||
|
||||
libero_env_instance = None
|
||||
task_desc = ""
|
||||
initial_states = None
|
||||
|
||||
max_infer_times = TASK_MAX_STEPS[args.task_suite_name]
|
||||
print(
|
||||
f"{args.task_suite_name} TASK_MAX_STEPS: {TASK_MAX_STEPS[args.task_suite_name]}"
|
||||
)
|
||||
for ep_idx in range(args.num_trials_per_task):
|
||||
print(f" > Running trial {ep_idx + 1} / {args.num_trials_per_task}...")
|
||||
|
||||
try:
|
||||
print(f"\nCreating environment for Task {task_id}...")
|
||||
libero_env_instance, task_desc, initial_states = (
|
||||
evaluator.create_env_for_task(task_id)
|
||||
)
|
||||
print(
|
||||
f"--- Starting task {task_id + 1} / {evaluator.num_tasks}: {task_desc} ---"
|
||||
)
|
||||
except Exception as e:
|
||||
print(
|
||||
f"\n[CRITICAL ERROR] Failed to create environment for task {task_id}: {e}. Skipping entire task."
|
||||
)
|
||||
continue
|
||||
|
||||
task_episodes_attempted += 1
|
||||
total_episodes_run += 1
|
||||
|
||||
success = False
|
||||
try:
|
||||
if args.mode == "flow":
|
||||
success = evaluator.run_infer_flow_action(
|
||||
env=libero_env_instance,
|
||||
task_id=task_id,
|
||||
task_desc=task_desc,
|
||||
default_initial_states=initial_states,
|
||||
episode_idx=ep_idx,
|
||||
max_infer_times=max_infer_times,
|
||||
)
|
||||
elif args.mode == "ar":
|
||||
success = evaluator.run_infer_ar_action(
|
||||
env=libero_env_instance,
|
||||
task_id=task_id,
|
||||
task_desc=task_desc,
|
||||
default_initial_states=initial_states,
|
||||
episode_idx=ep_idx,
|
||||
max_infer_times=max_infer_times,
|
||||
)
|
||||
except Exception as e:
|
||||
print(f" [EXCEPTION] Episode run error: {e}")
|
||||
|
||||
if success:
|
||||
task_successes += 1
|
||||
total_successes += 1
|
||||
print(" > Trial result: SUCCESS")
|
||||
else:
|
||||
print(" > Trial result: FAILURE")
|
||||
|
||||
if task_episodes_attempted > 0:
|
||||
print(
|
||||
f" > Task {task_id} current success rate: {task_successes / task_episodes_attempted * 100:.1f}% ({task_successes}/{task_episodes_attempted})"
|
||||
)
|
||||
if total_episodes_run > 0:
|
||||
print(
|
||||
f" > Overall current success rate: {total_successes / total_episodes_run * 100:.1f}% ({total_successes}/{total_episodes_run})"
|
||||
)
|
||||
|
||||
task_success_rate = (
|
||||
task_successes / task_episodes_attempted
|
||||
if task_episodes_attempted > 0
|
||||
else 0
|
||||
)
|
||||
print(f"\n--- Task {task_id} ({task_desc}) Summary ---")
|
||||
print(
|
||||
f"Success rate: {task_success_rate * 100:.1f}% ({task_successes}/{task_episodes_attempted})"
|
||||
)
|
||||
print(f"{'-'*40}\n")
|
||||
|
||||
end_time = time.time()
|
||||
total_time = end_time - start_time
|
||||
final_success_rate = (
|
||||
total_successes / total_episodes_run if total_episodes_run > 0 else 0
|
||||
)
|
||||
|
||||
print(f"\n{'='*20} Final Evaluation Summary {'='*20}")
|
||||
print(f"Total runtime: {total_time:.2f} seconds ({total_time / 60:.1f} minutes)")
|
||||
print(f"Total trials run: {total_episodes_run}")
|
||||
print(f"Total successes: {total_successes}")
|
||||
print(f"Overall success rate: {final_success_rate * 100:.2f}%")
|
||||
print(f"{'='*56}")
|
||||
|
||||
print("Evaluation completed.")
|
||||
raise SystemExit(main())
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -6,9 +6,10 @@ Works around the StorageMeta compatibility issue between PyTorch versions.
|
||||
|
||||
import os
|
||||
import sys
|
||||
import torch
|
||||
from pathlib import Path
|
||||
from typing import Dict
|
||||
|
||||
import torch
|
||||
from safetensors.torch import save_file
|
||||
|
||||
|
||||
@@ -70,65 +71,16 @@ def load_sharded_checkpoint(checkpoint_dir: str) -> Dict[str, torch.Tensor]:
|
||||
|
||||
except AttributeError as e:
|
||||
if "StorageMeta" in str(e):
|
||||
print(f"[ERROR] StorageMeta compatibility issue: {e}")
|
||||
print("[INFO] Attempting alternative loading method...")
|
||||
return load_checkpoint_alternative(checkpoint_dir)
|
||||
else:
|
||||
raise
|
||||
|
||||
|
||||
def load_checkpoint_alternative(checkpoint_dir: str) -> Dict[str, torch.Tensor]:
|
||||
"""
|
||||
Alternative method to load checkpoint by directly reading shard files.
|
||||
|
||||
Args:
|
||||
checkpoint_dir: Path to directory containing .distcp files
|
||||
|
||||
Returns:
|
||||
Dictionary of merged model state
|
||||
"""
|
||||
checkpoint_path = Path(checkpoint_dir)
|
||||
|
||||
# Find all shard files
|
||||
shard_files = sorted(checkpoint_path.glob("*.distcp"))
|
||||
|
||||
if not shard_files:
|
||||
raise FileNotFoundError(f"No .distcp files found in {checkpoint_dir}")
|
||||
|
||||
print(f"[INFO] Found {len(shard_files)} shard files")
|
||||
|
||||
# Load all shards
|
||||
merged_state = {}
|
||||
|
||||
for shard_file in shard_files:
|
||||
print(f"[INFO] Loading shard: {shard_file.name}")
|
||||
try:
|
||||
shard_data = torch.load(shard_file, map_location="cpu")
|
||||
|
||||
# Merge the shard into the state dict
|
||||
if isinstance(shard_data, dict):
|
||||
for key, value in shard_data.items():
|
||||
if isinstance(value, torch.Tensor):
|
||||
if key in merged_state:
|
||||
# Handle duplicates - concatenate or overwrite based on shape
|
||||
print(f"[WARNING] Duplicate key found: {key}")
|
||||
merged_state[key] = value
|
||||
elif isinstance(value, dict):
|
||||
# Nested dict structure
|
||||
for subkey, subvalue in value.items():
|
||||
full_key = f"{key}.{subkey}" if key else subkey
|
||||
if isinstance(subvalue, torch.Tensor):
|
||||
merged_state[full_key] = subvalue
|
||||
|
||||
except Exception as e:
|
||||
print(f"[WARNING] Failed to load shard {shard_file.name}: {e}")
|
||||
continue
|
||||
|
||||
if not merged_state:
|
||||
raise RuntimeError("Failed to load any checkpoint data from shards")
|
||||
|
||||
print(f"[INFO] Loaded {len(merged_state)} tensors from shards")
|
||||
return merged_state
|
||||
raise RuntimeError(
|
||||
"Unable to load this DCP checkpoint because its metadata uses "
|
||||
"StorageMeta from a different PyTorch version. The previous "
|
||||
"manual .distcp fallback was removed because FSDP shards cannot "
|
||||
"be reconstructed by directly loading shard files and overwriting "
|
||||
"duplicate keys. Please run this script with a PyTorch version "
|
||||
"compatible with the checkpoint writer, or re-save the checkpoint "
|
||||
"with the current PyTorch version."
|
||||
) from e
|
||||
raise
|
||||
|
||||
|
||||
def save_merged_checkpoint(
|
||||
|
||||
+101
-20
@@ -1,26 +1,107 @@
|
||||
from transformers import AutoProcessor
|
||||
import os
|
||||
#!/usr/bin/env python3
|
||||
"""Merge Wall-X action tokens into a Qwen2.5-VL processor tokenizer."""
|
||||
|
||||
processor_path = "/path/to/Qwen2.5-VL-3B-Instruct"
|
||||
action_tokenizer_path = "/path/to/fast"
|
||||
use_fast_tokenizer = True
|
||||
from __future__ import annotations
|
||||
|
||||
processor = AutoProcessor.from_pretrained(processor_path, use_fast=True)
|
||||
processor.tokenizer.padding_side = "left"
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
action_tokenizer = AutoProcessor.from_pretrained(
|
||||
action_tokenizer_path, trust_remote_code=True
|
||||
)
|
||||
|
||||
new_tokens = ["<|propri|>", "<|action|>"]
|
||||
new_tokens += [f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size)]
|
||||
num_added_tokens = processor.tokenizer.add_tokens(new_tokens)
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description=(
|
||||
"Create a Wall-X processor directory by adding FAST action tokens "
|
||||
"to a Qwen2.5-VL processor tokenizer."
|
||||
)
|
||||
)
|
||||
parser.add_argument(
|
||||
"--processor-path",
|
||||
required=True,
|
||||
help="Base Qwen2.5-VL processor directory or Hugging Face repo id.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--action-tokenizer-path",
|
||||
required=True,
|
||||
help="FAST/action tokenizer processor directory or Hugging Face repo id.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
required=True,
|
||||
help="Directory where the merged processor will be written.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use-fast",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="Use the fast tokenizer implementation when loading the base processor.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--trust-remote-code",
|
||||
action="store_true",
|
||||
help="Allow custom code when loading the action tokenizer processor.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
begin_idx_token = "<|action_token_0|>"
|
||||
token_id = processor.tokenizer.convert_tokens_to_ids(begin_idx_token)
|
||||
processor.tokenizer.init_kwargs["action_token_start_index"] = token_id
|
||||
processor.tokenizer.init_kwargs["action_token_vocab_size"] = action_tokenizer.vocab_size
|
||||
|
||||
new_tokenizer_dir = "/path/to/new_tokenizer"
|
||||
os.makedirs(new_tokenizer_dir, exist_ok=True)
|
||||
processor.save_pretrained(new_tokenizer_dir)
|
||||
def _resolve_action_vocab_size(action_processor) -> int:
|
||||
vocab_size = getattr(action_processor, "vocab_size", None)
|
||||
if vocab_size is None and hasattr(action_processor, "tokenizer"):
|
||||
vocab_size = getattr(action_processor.tokenizer, "vocab_size", None)
|
||||
if vocab_size is None:
|
||||
raise AttributeError(
|
||||
"Could not determine action tokenizer vocab size from the loaded processor."
|
||||
)
|
||||
return int(vocab_size)
|
||||
|
||||
|
||||
def merge_tokenizer(
|
||||
*,
|
||||
processor_path: str,
|
||||
action_tokenizer_path: str,
|
||||
output_dir: str,
|
||||
use_fast: bool,
|
||||
trust_remote_code: bool,
|
||||
) -> None:
|
||||
from transformers import AutoProcessor
|
||||
|
||||
processor = AutoProcessor.from_pretrained(processor_path, use_fast=use_fast)
|
||||
processor.tokenizer.padding_side = "left"
|
||||
|
||||
action_processor = AutoProcessor.from_pretrained(
|
||||
action_tokenizer_path,
|
||||
trust_remote_code=trust_remote_code,
|
||||
)
|
||||
action_vocab_size = _resolve_action_vocab_size(action_processor)
|
||||
|
||||
new_tokens = ["<|propri|>", "<|action|>"]
|
||||
new_tokens += [f"<|action_token_{i}|>" for i in range(action_vocab_size)]
|
||||
num_added_tokens = processor.tokenizer.add_tokens(new_tokens)
|
||||
|
||||
begin_idx_token = "<|action_token_0|>"
|
||||
token_id = processor.tokenizer.convert_tokens_to_ids(begin_idx_token)
|
||||
processor.tokenizer.init_kwargs["action_token_start_index"] = token_id
|
||||
processor.tokenizer.init_kwargs["action_token_vocab_size"] = action_vocab_size
|
||||
|
||||
output_path = Path(output_dir)
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
processor.save_pretrained(output_path)
|
||||
|
||||
print(f"Saved merged processor to {output_path}")
|
||||
print(f"Added {num_added_tokens} tokenizer tokens")
|
||||
print(f"action_token_start_index={token_id}")
|
||||
print(f"action_token_vocab_size={action_vocab_size}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
merge_tokenizer(
|
||||
processor_path=args.processor_path,
|
||||
action_tokenizer_path=args.action_tokenizer_path,
|
||||
output_dir=args.output_dir,
|
||||
use_fast=args.use_fast,
|
||||
trust_remote_code=args.trust_remote_code,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -1,161 +0,0 @@
|
||||
# This file is copied from openpi
|
||||
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())
|
||||
Executable
+273
@@ -0,0 +1,273 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
# LIBERO evaluation launcher. Configure with environment variables:
|
||||
#
|
||||
# CHECKPOINT_PATH=/path/to/checkpoint bash scripts/run_libero.sh
|
||||
# bash scripts/run_libero.sh /path/to/checkpoint
|
||||
# CHECKPOINT_PATH=/path/to/checkpoint SMOKE=1 bash scripts/run_libero.sh
|
||||
# CONFIG=/path/to/eval.yaml bash scripts/run_libero.sh
|
||||
#
|
||||
# Optional knobs:
|
||||
# LIBERO_PATH=/path/to/LIBERO # optional when not cloned to third_party/LIBERO
|
||||
# CUDA_ID=0
|
||||
# DRIVER_MODE=in_process
|
||||
# NUM_WORKERS=1
|
||||
# MAX_BATCH_SIZE=1
|
||||
# ALL_SUITES=1 # run all standard LIBERO suites (40 tasks)
|
||||
# TASK_SUITES="libero_spatial ..." # custom suite list (space- or comma-separated)
|
||||
# TASK_SUITE_NAME=libero_spatial # single suite when ALL_SUITES=0 and TASK_SUITES unset
|
||||
# NUM_TRIALS_PER_TASK=50
|
||||
# TASK_INDICES=0,1,2 # omit to run every task in the suite
|
||||
# MAX_INFER_TIMES=52
|
||||
# NORM_KEY=libero_all
|
||||
# ROLLOUT_BASE=/path/to/rollout # per-suite logs under ${ROLLOUT_BASE}/${suite}/
|
||||
# LOG_DIR=/tmp/harrix_libero_eval # overrides ROLLOUT_BASE when set (single suite)
|
||||
# SKIP_LIBERO_DEP_CHECK=1 # bypass dependency preflight
|
||||
|
||||
# export ALL_SUITES=1
|
||||
# export TASK_SUITE_NAME=libero_10
|
||||
# export CUDA_ID=0
|
||||
# export NUM_WORKERS=10
|
||||
# export NUM_TRIALS_PER_TASK=20
|
||||
# export CHECKPOINT_PATH=/path/to/checkpoint
|
||||
# export ROLLOUT_BASE=/path/to/rollout
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
SOURCE_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
|
||||
|
||||
if [[ -d "${SOURCE_ROOT}/third_party/LIBERO" ]]; then
|
||||
export PYTHONPATH="${SOURCE_ROOT}/third_party/LIBERO:${PYTHONPATH:-}"
|
||||
fi
|
||||
if [[ -n "${LIBERO_PATH:-}" ]]; then
|
||||
export PYTHONPATH="${LIBERO_PATH}:${PYTHONPATH:-}"
|
||||
fi
|
||||
|
||||
if [[ -f "${SCRIPT_DIR}/infer_libero.py" ]]; then
|
||||
INFER_LIBERO_CMD=("${SCRIPT_DIR}/infer_libero.py")
|
||||
elif [[ -f "${SOURCE_ROOT}/scripts/infer_libero.py" ]]; then
|
||||
INFER_LIBERO_CMD=(python "${SOURCE_ROOT}/scripts/infer_libero.py")
|
||||
elif command -v infer_libero.py >/dev/null 2>&1; then
|
||||
INFER_LIBERO_CMD=(infer_libero.py)
|
||||
else
|
||||
echo "infer_libero.py is not available. Install Wall-X first." >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
DEFAULT_ALL_SUITES=(
|
||||
libero_spatial
|
||||
libero_object
|
||||
libero_goal
|
||||
libero_10
|
||||
)
|
||||
|
||||
if [[ $# -gt 1 ]]; then
|
||||
echo "Usage: bash scripts/run_libero.sh [CHECKPOINT_PATH]" >&2
|
||||
exit 2
|
||||
fi
|
||||
if [[ $# -eq 1 ]]; then
|
||||
CHECKPOINT_PATH="$1"
|
||||
fi
|
||||
|
||||
export CUDA_VISIBLE_DEVICES="${CUDA_ID:-0}"
|
||||
|
||||
# MuJoCo offscreen rendering via NVIDIA EGL (required on headless GPU nodes).
|
||||
export MUJOCO_GL=egl
|
||||
export PYOPENGL_PLATFORM=egl
|
||||
# LIBERO init-state files are trusted simulator assets. PyTorch 2.6 changed
|
||||
# torch.load() defaults in a way that breaks LIBERO's upstream loader.
|
||||
export TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD="${TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD:-1}"
|
||||
EGL_VENDOR_DIR="${HOME}/.config/glvnd/egl_vendor.d"
|
||||
EGL_VENDOR_FILE="${EGL_VENDOR_DIR}/10_nvidia.json"
|
||||
if [[ ! -f "${EGL_VENDOR_FILE}" ]]; then
|
||||
mkdir -p "${EGL_VENDOR_DIR}"
|
||||
cat > "${EGL_VENDOR_FILE}" <<'EOF'
|
||||
{
|
||||
"file_format_version" : "1.0.0",
|
||||
"ICD" : {
|
||||
"library_path" : "libEGL_nvidia.so.0"
|
||||
}
|
||||
}
|
||||
EOF
|
||||
fi
|
||||
export __EGL_VENDOR_LIBRARY_FILENAMES="${EGL_VENDOR_FILE}"
|
||||
export LD_LIBRARY_PATH="/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib:/usr/local/nvidia/lib64:${LD_LIBRARY_PATH:-}"
|
||||
|
||||
# After CUDA_VISIBLE_DEVICES remapping, MuJoCo only sees devices from index 0.
|
||||
export MUJOCO_EGL_DEVICE_ID="${MUJOCO_EGL_DEVICE_ID:-0}"
|
||||
if [[ -d "${SOURCE_ROOT}/third_party/harrix/python" ]]; then
|
||||
export PYTHONPATH="${SOURCE_ROOT}/third_party/harrix/python:${PYTHONPATH:-}"
|
||||
fi
|
||||
if [[ -d "${SOURCE_ROOT}/wall_x" ]]; then
|
||||
export PYTHONPATH="${SOURCE_ROOT}:${PYTHONPATH:-}"
|
||||
fi
|
||||
|
||||
check_libero_dependencies() {
|
||||
local missing
|
||||
if missing="$(
|
||||
python - <<'PY'
|
||||
import importlib.util
|
||||
|
||||
checks = [
|
||||
("libero.libero", "LIBERO"),
|
||||
("robosuite", "robosuite"),
|
||||
("mujoco", "mujoco"),
|
||||
("OpenGL", "PyOpenGL"),
|
||||
("bddl", "bddl"),
|
||||
("gym", "gym"),
|
||||
("h5py", "h5py"),
|
||||
]
|
||||
|
||||
missing = [label for module, label in checks if importlib.util.find_spec(module) is None]
|
||||
if missing:
|
||||
print(", ".join(missing))
|
||||
raise SystemExit(1)
|
||||
PY
|
||||
)"; then
|
||||
return 0
|
||||
fi
|
||||
|
||||
cat >&2 <<EOF
|
||||
LIBERO simulator dependencies are missing: ${missing}
|
||||
|
||||
Install the optional LIBERO simulator stack from the Wall-X repository root:
|
||||
|
||||
pip install -r requirements-libero.txt
|
||||
mkdir -p third_party
|
||||
git clone https://github.com/Lifelong-Robot-Learning/LIBERO.git third_party/LIBERO
|
||||
|
||||
If LIBERO is cloned elsewhere, pass it with:
|
||||
|
||||
LIBERO_PATH=/path/to/LIBERO bash scripts/run_libero.sh ...
|
||||
|
||||
Set SKIP_LIBERO_DEP_CHECK=1 only if you intentionally manage these dependencies elsewhere.
|
||||
EOF
|
||||
exit 2
|
||||
}
|
||||
|
||||
if [[ "${SKIP_LIBERO_DEP_CHECK:-0}" != "1" ]]; then
|
||||
check_libero_dependencies
|
||||
fi
|
||||
|
||||
resolve_task_suites() {
|
||||
local suites=()
|
||||
if [[ "${ALL_SUITES:-0}" == "1" ]]; then
|
||||
if [[ -n "${TASK_SUITES:-}" ]]; then
|
||||
TASK_SUITES="${TASK_SUITES//,/ }"
|
||||
read -r -a suites <<< "${TASK_SUITES}"
|
||||
else
|
||||
suites=("${DEFAULT_ALL_SUITES[@]}")
|
||||
fi
|
||||
elif [[ -n "${TASK_SUITES:-}" ]]; then
|
||||
TASK_SUITES="${TASK_SUITES//,/ }"
|
||||
read -r -a suites <<< "${TASK_SUITES}"
|
||||
else
|
||||
suites=("${TASK_SUITE_NAME:-libero_spatial}")
|
||||
fi
|
||||
echo "${suites[@]}"
|
||||
}
|
||||
|
||||
resolve_log_dirs() {
|
||||
local suite="$1"
|
||||
local multi_suite="$2"
|
||||
local log_dir rollout_dir
|
||||
|
||||
if [[ -n "${LOG_DIR:-}" ]]; then
|
||||
if [[ "${multi_suite}" == "1" ]]; then
|
||||
log_dir="${LOG_DIR}/${suite}"
|
||||
else
|
||||
log_dir="${LOG_DIR}"
|
||||
fi
|
||||
elif [[ -n "${ROLLOUT_BASE:-}" ]]; then
|
||||
if [[ "${multi_suite}" == "1" ]]; then
|
||||
log_dir="${ROLLOUT_BASE}/${suite}/json"
|
||||
else
|
||||
log_dir="${ROLLOUT_BASE}/json"
|
||||
fi
|
||||
else
|
||||
log_dir="/tmp/harrix_libero_eval"
|
||||
fi
|
||||
|
||||
if [[ -n "${WALLX_ROLLOUT_DIR:-}" ]]; then
|
||||
if [[ "${multi_suite}" == "1" ]]; then
|
||||
rollout_dir="${WALLX_ROLLOUT_DIR}/${suite}"
|
||||
else
|
||||
rollout_dir="${WALLX_ROLLOUT_DIR}"
|
||||
fi
|
||||
elif [[ -n "${ROLLOUT_BASE:-}" ]]; then
|
||||
if [[ "${multi_suite}" == "1" ]]; then
|
||||
rollout_dir="${ROLLOUT_BASE}/${suite}/videos"
|
||||
else
|
||||
rollout_dir="${ROLLOUT_BASE}/videos"
|
||||
fi
|
||||
else
|
||||
rollout_dir=""
|
||||
fi
|
||||
|
||||
mkdir -p "${log_dir}"
|
||||
if [[ -n "${rollout_dir}" ]]; then
|
||||
mkdir -p "${rollout_dir}"
|
||||
fi
|
||||
LOG_DIR_RESOLVED="${log_dir}"
|
||||
WALLX_ROLLOUT_DIR_RESOLVED="${rollout_dir}"
|
||||
}
|
||||
|
||||
run_suite() {
|
||||
local suite="$1"
|
||||
local multi_suite="$2"
|
||||
|
||||
resolve_log_dirs "${suite}" "${multi_suite}"
|
||||
if [[ -n "${WALLX_ROLLOUT_DIR_RESOLVED}" ]]; then
|
||||
export WALLX_ROLLOUT_DIR="${WALLX_ROLLOUT_DIR_RESOLVED}"
|
||||
else
|
||||
unset WALLX_ROLLOUT_DIR
|
||||
fi
|
||||
|
||||
local args=()
|
||||
if [[ -n "${CONFIG:-}" ]]; then
|
||||
args+=(--config "${CONFIG}")
|
||||
else
|
||||
if [[ -z "${CHECKPOINT_PATH:-}" ]]; then
|
||||
echo "CHECKPOINT_PATH is required when CONFIG is not set." >&2
|
||||
exit 2
|
||||
fi
|
||||
args+=(--checkpoint-path "${CHECKPOINT_PATH}")
|
||||
fi
|
||||
|
||||
if [[ -n "${TRAIN_CONFIG_PATH:-}" ]]; then
|
||||
args+=(--train-config-path "${TRAIN_CONFIG_PATH}")
|
||||
fi
|
||||
if [[ "${SMOKE:-0}" == "1" ]]; then
|
||||
args+=(--smoke)
|
||||
fi
|
||||
if [[ "${DET:-0}" == "1" ]]; then
|
||||
args+=(--deterministic-model)
|
||||
fi
|
||||
if [[ -n "${TASK_INDICES:-}" ]]; then
|
||||
args+=(--task-indices "${TASK_INDICES}")
|
||||
fi
|
||||
if [[ -n "${MAX_INFER_TIMES:-}" ]]; then
|
||||
args+=(--max-infer-times "${MAX_INFER_TIMES}")
|
||||
fi
|
||||
|
||||
args+=(--driver-mode "${DRIVER_MODE:-in_process}")
|
||||
args+=(--num-workers "${NUM_WORKERS:-1}")
|
||||
args+=(--max-batch-size "${MAX_BATCH_SIZE:-1}")
|
||||
args+=(--task-suite-name "${suite}")
|
||||
args+=(--num-trials-per-task "${NUM_TRIALS_PER_TASK:-50}")
|
||||
args+=(--norm-key "${NORM_KEY:-libero_all}")
|
||||
args+=(--log-dir "${LOG_DIR_RESOLVED}")
|
||||
|
||||
echo "=== LIBERO eval: suite=${suite} log_dir=${LOG_DIR_RESOLVED} ==="
|
||||
"${INFER_LIBERO_CMD[@]}" "${args[@]}"
|
||||
}
|
||||
|
||||
suites=($(resolve_task_suites))
|
||||
multi_suite=0
|
||||
if [[ "${#suites[@]}" -gt 1 ]]; then
|
||||
multi_suite=1
|
||||
fi
|
||||
|
||||
for suite in "${suites[@]}"; do
|
||||
run_suite "${suite}" "${multi_suite}"
|
||||
done
|
||||
Executable
+285
@@ -0,0 +1,285 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
usage() {
|
||||
cat <<'EOF'
|
||||
Usage:
|
||||
bash scripts/run_serving.sh --checkpoint-path /path/to/checkpoint [options]
|
||||
|
||||
Required:
|
||||
--checkpoint-path PATH Checkpoint directory or checkpoint file.
|
||||
|
||||
Common options:
|
||||
--train-config-path PATH Training config used by the checkpoint.
|
||||
--port PORT WebSocket port. Default: 32195.
|
||||
--host HOST Bind host. Default: 0.0.0.0.
|
||||
--env X2ROBOT|LIBERO Serving environment. Default: X2ROBOT.
|
||||
--cuda-id ID Sets CUDA_VISIBLE_DEVICES. Default: 0.
|
||||
--image-passing-mode MODE base64 or numpy. Default: base64.
|
||||
--action-horizon N Model action horizon. Default: 32.
|
||||
--robot-type TYPE desktop, turtle, or ex001. Default: desktop.
|
||||
--raw-actions Alias for --no-serialize-actions.
|
||||
--serialize-actions Return robot-serialized actions.
|
||||
--no-serialize-actions Return raw model action chunks. Default.
|
||||
--max-batch-size N Enable dynamic batching.
|
||||
--enable-cuda-graph Enable CUDA graph in the serving runtime.
|
||||
--enable-experimental-engine Enable the experimental inference engine.
|
||||
--debug Enable debug logging.
|
||||
--dry-run Print the command without running it.
|
||||
|
||||
Additional arguments after "--" are forwarded to launch_serving.py, for example:
|
||||
bash scripts/run_serving.sh --checkpoint-path /ckpt -- \
|
||||
--model-config.norm-key libero_all
|
||||
|
||||
Environment variables can also be used, e.g. CHECKPOINT_PATH,
|
||||
TRAIN_CONFIG_PATH, PORT, CUDA_ID, ACTION_HORIZON, WALLX_ENV.
|
||||
EOF
|
||||
}
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
SOURCE_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
|
||||
|
||||
if [[ -d "${SOURCE_ROOT}/wall_x" ]]; then
|
||||
export PYTHONPATH="${SOURCE_ROOT}:${PYTHONPATH:-}"
|
||||
fi
|
||||
|
||||
PYTHON_BIN="${PYTHON_BIN:-python}"
|
||||
CHECKPOINT_PATH="${CHECKPOINT_PATH:-}"
|
||||
TRAIN_CONFIG_PATH="${TRAIN_CONFIG_PATH:-}"
|
||||
PORT="${PORT:-32195}"
|
||||
HOST="${HOST:-0.0.0.0}"
|
||||
WALLX_ENV="${WALLX_ENV:-X2ROBOT}"
|
||||
CUDA_ID="${CUDA_ID:-0}"
|
||||
IMAGE_PASSING_MODE="${IMAGE_PASSING_MODE:-base64}"
|
||||
ACTION_HORIZON="${ACTION_HORIZON:-32}"
|
||||
ROBOT_TYPE="${ROBOT_TYPE:-desktop}"
|
||||
ROBOT_ACTION_INTERPOLATE_MULTIPLIER="${ROBOT_ACTION_INTERPOLATE_MULTIPLIER:-1}"
|
||||
ROBOT_ACTION_END_RATIO="${ROBOT_ACTION_END_RATIO:-1.0}"
|
||||
MODEL_DEVICE="${MODEL_DEVICE:-cuda}"
|
||||
MAX_BATCH_SIZE="${MAX_BATCH_SIZE:-}"
|
||||
DEFAULT_PROMPT="${DEFAULT_PROMPT:-}"
|
||||
SERIALIZE_ACTIONS="${SERIALIZE_ACTIONS:-0}"
|
||||
ENABLE_CUDA_GRAPH="${ENABLE_CUDA_GRAPH:-0}"
|
||||
ENABLE_EXPERIMENTAL_ENGINE="${ENABLE_EXPERIMENTAL_ENGINE:-0}"
|
||||
DEBUG="${DEBUG:-0}"
|
||||
DRY_RUN=0
|
||||
EXTRA_ARGS=()
|
||||
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case "$1" in
|
||||
-h|--help)
|
||||
usage
|
||||
exit 0
|
||||
;;
|
||||
--checkpoint-path)
|
||||
CHECKPOINT_PATH="${2:?missing value for --checkpoint-path}"
|
||||
shift 2
|
||||
;;
|
||||
--checkpoint-path=*)
|
||||
CHECKPOINT_PATH="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--train-config-path)
|
||||
TRAIN_CONFIG_PATH="${2:?missing value for --train-config-path}"
|
||||
shift 2
|
||||
;;
|
||||
--train-config-path=*)
|
||||
TRAIN_CONFIG_PATH="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--port)
|
||||
PORT="${2:?missing value for --port}"
|
||||
shift 2
|
||||
;;
|
||||
--port=*)
|
||||
PORT="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--host)
|
||||
HOST="${2:?missing value for --host}"
|
||||
shift 2
|
||||
;;
|
||||
--host=*)
|
||||
HOST="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--env)
|
||||
WALLX_ENV="${2:?missing value for --env}"
|
||||
shift 2
|
||||
;;
|
||||
--env=*)
|
||||
WALLX_ENV="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--cuda-id)
|
||||
CUDA_ID="${2:?missing value for --cuda-id}"
|
||||
shift 2
|
||||
;;
|
||||
--cuda-id=*)
|
||||
CUDA_ID="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--image-passing-mode)
|
||||
IMAGE_PASSING_MODE="${2:?missing value for --image-passing-mode}"
|
||||
shift 2
|
||||
;;
|
||||
--image-passing-mode=*)
|
||||
IMAGE_PASSING_MODE="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--action-horizon)
|
||||
ACTION_HORIZON="${2:?missing value for --action-horizon}"
|
||||
shift 2
|
||||
;;
|
||||
--action-horizon=*)
|
||||
ACTION_HORIZON="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--robot-type)
|
||||
ROBOT_TYPE="${2:?missing value for --robot-type}"
|
||||
shift 2
|
||||
;;
|
||||
--robot-type=*)
|
||||
ROBOT_TYPE="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--robot-action-interpolate-multiplier)
|
||||
ROBOT_ACTION_INTERPOLATE_MULTIPLIER="${2:?missing value for --robot-action-interpolate-multiplier}"
|
||||
shift 2
|
||||
;;
|
||||
--robot-action-interpolate-multiplier=*)
|
||||
ROBOT_ACTION_INTERPOLATE_MULTIPLIER="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--robot-action-end-ratio)
|
||||
ROBOT_ACTION_END_RATIO="${2:?missing value for --robot-action-end-ratio}"
|
||||
shift 2
|
||||
;;
|
||||
--robot-action-end-ratio=*)
|
||||
ROBOT_ACTION_END_RATIO="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--model-device)
|
||||
MODEL_DEVICE="${2:?missing value for --model-device}"
|
||||
shift 2
|
||||
;;
|
||||
--model-device=*)
|
||||
MODEL_DEVICE="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--max-batch-size)
|
||||
MAX_BATCH_SIZE="${2:?missing value for --max-batch-size}"
|
||||
shift 2
|
||||
;;
|
||||
--max-batch-size=*)
|
||||
MAX_BATCH_SIZE="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--default-prompt)
|
||||
DEFAULT_PROMPT="${2:?missing value for --default-prompt}"
|
||||
shift 2
|
||||
;;
|
||||
--default-prompt=*)
|
||||
DEFAULT_PROMPT="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--serialize-actions)
|
||||
SERIALIZE_ACTIONS=1
|
||||
shift
|
||||
;;
|
||||
--no-serialize-actions|--raw-actions)
|
||||
SERIALIZE_ACTIONS=0
|
||||
shift
|
||||
;;
|
||||
--enable-cuda-graph)
|
||||
ENABLE_CUDA_GRAPH=1
|
||||
shift
|
||||
;;
|
||||
--enable-experimental-engine)
|
||||
ENABLE_EXPERIMENTAL_ENGINE=1
|
||||
shift
|
||||
;;
|
||||
--debug)
|
||||
DEBUG=1
|
||||
shift
|
||||
;;
|
||||
--dry-run)
|
||||
DRY_RUN=1
|
||||
shift
|
||||
;;
|
||||
--)
|
||||
shift
|
||||
EXTRA_ARGS+=("$@")
|
||||
break
|
||||
;;
|
||||
*)
|
||||
EXTRA_ARGS+=("$1")
|
||||
shift
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
if [[ -z "${CHECKPOINT_PATH}" ]]; then
|
||||
echo "error: --checkpoint-path is required." >&2
|
||||
echo >&2
|
||||
usage >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
export CUDA_VISIBLE_DEVICES="${CUDA_ID}"
|
||||
export ENABLE_FAST_PREPROCESS="${ENABLE_FAST_PREPROCESS:-true}"
|
||||
|
||||
CMD=(
|
||||
"${PYTHON_BIN}" -m wall_x._vendor.harrix.serving.launch_serving
|
||||
--env "${WALLX_ENV}"
|
||||
--host "${HOST}"
|
||||
--port "${PORT}"
|
||||
--image-passing-mode "${IMAGE_PASSING_MODE}"
|
||||
)
|
||||
|
||||
if [[ "${SERIALIZE_ACTIONS}" == "1" ]]; then
|
||||
CMD+=(--serialize-actions)
|
||||
else
|
||||
CMD+=(--no-serialize-actions)
|
||||
fi
|
||||
if [[ -n "${MAX_BATCH_SIZE}" ]]; then
|
||||
CMD+=(--max-batch-size "${MAX_BATCH_SIZE}")
|
||||
fi
|
||||
if [[ -n "${DEFAULT_PROMPT}" ]]; then
|
||||
CMD+=(--default-prompt "${DEFAULT_PROMPT}")
|
||||
fi
|
||||
if [[ "${ENABLE_CUDA_GRAPH}" == "1" ]]; then
|
||||
CMD+=(--enable-cuda-graph)
|
||||
fi
|
||||
if [[ "${ENABLE_EXPERIMENTAL_ENGINE}" == "1" ]]; then
|
||||
CMD+=(--enable-experimental-engine)
|
||||
fi
|
||||
if [[ "${DEBUG}" == "1" ]]; then
|
||||
CMD+=(--debug)
|
||||
fi
|
||||
|
||||
CMD+=(
|
||||
model-config:server-model-config
|
||||
--model-config.checkpoint-path "${CHECKPOINT_PATH}"
|
||||
--model-config.action-horizon "${ACTION_HORIZON}"
|
||||
--model-config.robot-type "${ROBOT_TYPE}"
|
||||
--model-config.robot-action-interpolate-multiplier "${ROBOT_ACTION_INTERPOLATE_MULTIPLIER}"
|
||||
--model-config.robot-action-end-ratio "${ROBOT_ACTION_END_RATIO}"
|
||||
--model-config.model-device "${MODEL_DEVICE}"
|
||||
)
|
||||
|
||||
if [[ -n "${TRAIN_CONFIG_PATH}" ]]; then
|
||||
CMD+=(--model-config.train-config-path "${TRAIN_CONFIG_PATH}")
|
||||
fi
|
||||
|
||||
CMD+=("${EXTRA_ARGS[@]}")
|
||||
|
||||
printf 'Launching Wall-X serving:\n'
|
||||
printf ' %q' "${CMD[@]}"
|
||||
printf '\n'
|
||||
|
||||
if [[ "${DRY_RUN}" == "1" ]]; then
|
||||
exit 0
|
||||
fi
|
||||
|
||||
exec "${CMD[@]}"
|
||||
@@ -1,104 +0,0 @@
|
||||
import torch
|
||||
from PIL import Image
|
||||
from transformers import AutoProcessor
|
||||
import yaml
|
||||
import os
|
||||
|
||||
from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction
|
||||
|
||||
|
||||
class VQAWrapper(object):
|
||||
def __init__(self, model_path: str, train_config: dict = None):
|
||||
|
||||
self.device = self._setup_device()
|
||||
if train_config is None:
|
||||
try:
|
||||
with open(os.path.join(model_path, "config.yml"), "r") as f:
|
||||
train_config = yaml.load(f, Loader=yaml.FullLoader)
|
||||
except Exception as e:
|
||||
print(f"load train_config.yml fail: {e}")
|
||||
self.processor = self._load_processor(train_config["processor_path"])
|
||||
self.model = self._load_model(model_path, train_config)
|
||||
|
||||
def _setup_device(self) -> str:
|
||||
if torch.cuda.is_available():
|
||||
return "cuda"
|
||||
else:
|
||||
return "cpu"
|
||||
|
||||
def _load_processor(self, model_path: str) -> AutoProcessor:
|
||||
return AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
|
||||
|
||||
def _load_model(
|
||||
self, model_path: str, train_config: dict
|
||||
) -> Qwen2_5_VLMoEForAction:
|
||||
model = Qwen2_5_VLMoEForAction.from_pretrained(
|
||||
model_path, train_config=train_config
|
||||
)
|
||||
if self.device == "cuda":
|
||||
model = model.to(self.device, dtype=torch.bfloat16)
|
||||
else:
|
||||
model.to(self.device)
|
||||
model.eval()
|
||||
return model
|
||||
|
||||
def generate(self, image: Image.Image, text: str, **kwargs) -> str:
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "image"}, {"type": "text", "text": text}],
|
||||
}
|
||||
]
|
||||
text_prompt = self.processor.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
inputs = self.processor(text=[text_prompt], images=[image], return_tensors="pt")
|
||||
inputs = {k: v.to(self.device) for k, v in inputs.items()}
|
||||
|
||||
generation_params = {
|
||||
"max_new_tokens": 1024, # default value, can be overridden by kwargs
|
||||
"do_sample": False,
|
||||
"eos_token_id": self.processor.tokenizer.eos_token_id,
|
||||
"pad_token_id": self.processor.tokenizer.pad_token_id,
|
||||
**kwargs,
|
||||
}
|
||||
|
||||
with torch.no_grad():
|
||||
generated_ids = self.model.generate(**inputs, **generation_params)
|
||||
|
||||
generated_ids = [
|
||||
output_ids[len(input_ids) :]
|
||||
for input_ids, output_ids in zip(inputs["input_ids"], generated_ids)
|
||||
]
|
||||
response = self.processor.batch_decode(
|
||||
generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
|
||||
)[0]
|
||||
return response
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
MODEL_PATH_FOR_MODULE_TEST = "/path/to/model_path"
|
||||
train_config_path = "/path/to/model_path/config.yml"
|
||||
with open(train_config_path, "r") as f:
|
||||
train_config = yaml.load(f, Loader=yaml.FullLoader)
|
||||
wrapper = VQAWrapper(
|
||||
model_path=MODEL_PATH_FOR_MODULE_TEST, train_config=train_config
|
||||
)
|
||||
|
||||
try:
|
||||
test_question = "To move the red block in the plate with same color, what should you do next? Think step by step."
|
||||
|
||||
# Local Image
|
||||
img = Image.open(
|
||||
"/x2robot_v2/yangping/github/wall-x/assets/cot_example_frame.png"
|
||||
).convert("RGB")
|
||||
# Internet Image
|
||||
# import requests
|
||||
# test_image_url = "https://www.ilankelman.org/stopsigns/australia.jpg"
|
||||
# img = Image.open(requests.get(test_image_url, stream=True).raw).convert("RGB")
|
||||
|
||||
answer = wrapper.generate(img, test_question)
|
||||
|
||||
print("model answer:", answer)
|
||||
except Exception as e:
|
||||
print(f"model answer fail: {e}")
|
||||
Reference in New Issue
Block a user