Update Wall-X to 1.1.0 (#104)
This commit is contained in:
Starrick Liu
@@ -0,0 +1,371 @@
|
||||
"""Load checkpoint-side train config and construct runtime config objects.
|
||||
|
||||
Inference adapters use this module to load the train config, apply checkpoint
|
||||
overrides for moved processor/tokenizer files, build model config, and build the
|
||||
data config required by preprocessing.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import logging
|
||||
|
||||
import yaml
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
_TASK_INFERENCE_KEYS = (
|
||||
"dof_config",
|
||||
"agent_pos_config",
|
||||
"ar_dof_config",
|
||||
"action_horizon",
|
||||
"action_horizon_flow",
|
||||
"noise_scheduler",
|
||||
"use_state_string_representation",
|
||||
)
|
||||
|
||||
_VIRTUAL_TAIL_KEYS = frozenset(("action_padding",))
|
||||
|
||||
|
||||
def _move_virtual_keys_to_tail(layout: dict | None) -> dict | None:
|
||||
"""Keep real LeRobot dims before virtual padding dims in inference layouts."""
|
||||
if not isinstance(layout, dict):
|
||||
return layout
|
||||
head = {k: v for k, v in layout.items() if k not in _VIRTUAL_TAIL_KEYS}
|
||||
tail = {k: v for k, v in layout.items() if k in _VIRTUAL_TAIL_KEYS}
|
||||
if not tail:
|
||||
return layout
|
||||
return {**head, **tail}
|
||||
|
||||
|
||||
def _canonicalize_task_layouts(target: dict) -> None:
|
||||
if os.environ.get("WALLX_CANONICALIZE_VIRTUAL_DOF_ORDER", "1") == "0":
|
||||
return
|
||||
task = target.get("task")
|
||||
data = target.get("data")
|
||||
for key in ("dof_config", "agent_pos_config", "ar_dof_config"):
|
||||
value = _move_virtual_keys_to_tail(target.get(key))
|
||||
if value is not None:
|
||||
target[key] = value
|
||||
if isinstance(task, dict):
|
||||
task_value = _move_virtual_keys_to_tail(task.get(key))
|
||||
if task_value is not None:
|
||||
task[key] = task_value
|
||||
if isinstance(data, dict):
|
||||
data_value = _move_virtual_keys_to_tail(data.get(key))
|
||||
if data_value is not None:
|
||||
data[key] = data_value
|
||||
|
||||
|
||||
def resolve_use_state_string_representation(train_config: dict) -> bool:
|
||||
"""Read ``use_state_string_representation`` with task YAML as authority."""
|
||||
task = train_config.get("task") or {}
|
||||
if task.get("use_state_string_representation") is not None:
|
||||
return bool(task["use_state_string_representation"])
|
||||
if train_config.get("use_state_string_representation") is not None:
|
||||
return bool(train_config["use_state_string_representation"])
|
||||
data = train_config.get("data") or {}
|
||||
if data.get("use_state_string_representation") is not None:
|
||||
return bool(data["use_state_string_representation"])
|
||||
return False
|
||||
|
||||
|
||||
def resolve_state_bins(train_config: dict, default: int = 256) -> int:
|
||||
"""Read discretization bin count from flat or nested train config."""
|
||||
if train_config.get("state_bins") is not None:
|
||||
return int(train_config["state_bins"])
|
||||
data = train_config.get("data") or {}
|
||||
if data.get("state_bins") is not None:
|
||||
return int(data["state_bins"])
|
||||
return default
|
||||
|
||||
|
||||
def resolve_agent_pos_config(train_config: dict) -> dict:
|
||||
"""Return ``agent_pos_config`` from flat or ``task``-nested train YAML."""
|
||||
task = train_config.get("task") or {}
|
||||
agent_pos_config = train_config.get("agent_pos_config") or task.get(
|
||||
"agent_pos_config"
|
||||
)
|
||||
if not agent_pos_config:
|
||||
raise KeyError(
|
||||
"agent_pos_config missing from train config "
|
||||
"(expected top-level or task.agent_pos_config)"
|
||||
)
|
||||
return dict(agent_pos_config)
|
||||
|
||||
|
||||
def resolve_dof_config(train_config: dict) -> dict:
|
||||
"""Return ``dof_config`` from flat or ``task``-nested train YAML."""
|
||||
task = train_config.get("task") or {}
|
||||
dof_config = train_config.get("dof_config") or task.get("dof_config")
|
||||
if not dof_config:
|
||||
raise KeyError(
|
||||
"dof_config missing from train config (expected top-level or task.dof_config)"
|
||||
)
|
||||
return dict(dof_config)
|
||||
|
||||
|
||||
def resolve_cam_names_from_train_config(train_config: dict) -> list[str] | None:
|
||||
"""Infer model camera keys from ``data.key_mappings.camera`` (e.g. LIBERO)."""
|
||||
data = train_config.get("data") or {}
|
||||
key_mappings = data.get("key_mappings") or {}
|
||||
camera = key_mappings.get("camera") or {}
|
||||
if not camera:
|
||||
return None
|
||||
names: list[str] = []
|
||||
for value in camera.values():
|
||||
name = str(value)
|
||||
if name not in names:
|
||||
names.append(name)
|
||||
return names or None
|
||||
|
||||
|
||||
def resolve_camera_label(cam_name: str, camera_name_mapping: dict | None = None) -> str:
|
||||
"""Match training ``get_wallx_normal_text`` camera display names."""
|
||||
mapping = camera_name_mapping or {}
|
||||
return mapping.get(cam_name, cam_name)
|
||||
|
||||
|
||||
def resolve_max_length(train_config: dict, default: int = 768) -> int:
|
||||
"""Read ``max_length`` from flat or nested train YAML (collator default 768)."""
|
||||
data = train_config.get("data") or {}
|
||||
if train_config.get("max_length") is not None:
|
||||
return int(train_config["max_length"])
|
||||
if data.get("max_length") is not None:
|
||||
return int(data["max_length"])
|
||||
raw = train_config.get("_raw_data") or {}
|
||||
if raw.get("max_length") is not None:
|
||||
return int(raw["max_length"])
|
||||
return default
|
||||
|
||||
|
||||
# Fields ``load_wallx_processors`` and model construction read from a flat dict.
|
||||
_ACTION_TOKENIZER_KEYS = (
|
||||
"action_tokenizer_type",
|
||||
"action_tokenizer_path",
|
||||
"action_tokenizer_checkpoint_path",
|
||||
"action_tokenizer_config_dir",
|
||||
"action_tokenizer",
|
||||
)
|
||||
|
||||
_INFERENCE_MODEL_KEYS = (
|
||||
"processor_path",
|
||||
"pretrained_path",
|
||||
"config_path",
|
||||
*_ACTION_TOKENIZER_KEYS,
|
||||
"ar_loss_weight",
|
||||
"attn_deterministic",
|
||||
"flow_loss_weight",
|
||||
"use_ema",
|
||||
)
|
||||
|
||||
|
||||
def _merge_model_fields(target: dict, source: dict, *, overwrite: bool = False) -> None:
|
||||
"""Copy model-related keys from ``source`` into flat ``target``."""
|
||||
if not source:
|
||||
return
|
||||
for key in _INFERENCE_MODEL_KEYS:
|
||||
value = source.get(key)
|
||||
if value is None:
|
||||
continue
|
||||
if overwrite or target.get(key) in (None, ""):
|
||||
target[key] = value
|
||||
|
||||
|
||||
def _mirror_task_fields(target: dict, task: dict) -> None:
|
||||
"""Copy typed task fields into the flat/data legacy inference mirrors."""
|
||||
if not isinstance(task, dict):
|
||||
return
|
||||
data = target.setdefault("data", {})
|
||||
if not isinstance(data, dict):
|
||||
data = {}
|
||||
target["data"] = data
|
||||
for key in _TASK_INFERENCE_KEYS:
|
||||
value = task.get(key)
|
||||
if value is None:
|
||||
continue
|
||||
target[key] = value
|
||||
data[key] = value
|
||||
|
||||
|
||||
def _load_ckpt_config_overlay(checkpoint_path: str) -> dict:
|
||||
"""Read ``config.yml`` saved beside a training checkpoint."""
|
||||
ckpt_yml = os.path.join(checkpoint_path, "config.yml")
|
||||
if not os.path.exists(ckpt_yml):
|
||||
return {}
|
||||
with open(ckpt_yml, "r") as f:
|
||||
raw = yaml.load(f, Loader=yaml.FullLoader) or {}
|
||||
|
||||
overlay: dict = {}
|
||||
_merge_model_fields(overlay, raw.get("model") or {}, overwrite=True)
|
||||
raw_yaml = raw.get("_raw_yaml") or {}
|
||||
_merge_model_fields(overlay, raw_yaml.get("model") or {}, overwrite=False)
|
||||
_merge_model_fields(overlay, raw, overwrite=False)
|
||||
return overlay
|
||||
|
||||
|
||||
def strip_action_tokenizer_fields(train_config: dict) -> None:
|
||||
"""Drop action-tokenizer fields so flow inference can skip AR tokenizer setup."""
|
||||
for key in _ACTION_TOKENIZER_KEYS:
|
||||
train_config.pop(key, None)
|
||||
model = train_config.get("model")
|
||||
if isinstance(model, dict):
|
||||
for key in _ACTION_TOKENIZER_KEYS:
|
||||
model.pop(key, None)
|
||||
|
||||
|
||||
def load_train_config_with_ckpt_overlay(
|
||||
train_config_path: str,
|
||||
checkpoint_path: str,
|
||||
) -> dict:
|
||||
"""Load train YAML and apply checkpoint-local processor/tokenizer overlays.
|
||||
|
||||
This lets a checkpoint remain portable when the original training-machine
|
||||
processor or action-tokenizer paths are no longer available.
|
||||
"""
|
||||
with open(train_config_path, "r") as f:
|
||||
train_config = yaml.load(f, Loader=yaml.FullLoader)
|
||||
|
||||
preprocessor_file = os.path.join(checkpoint_path, "preprocessor_config.json")
|
||||
if os.path.exists(preprocessor_file):
|
||||
train_config["processor_path"] = checkpoint_path
|
||||
|
||||
orig_action_tok = train_config.get("action_tokenizer_path", None)
|
||||
if orig_action_tok is not None and not os.path.exists(orig_action_tok):
|
||||
tokenizer_file = os.path.join(checkpoint_path, "tokenizer.json")
|
||||
tokenizer_config_file = os.path.join(checkpoint_path, "tokenizer_config.json")
|
||||
if os.path.exists(tokenizer_file) and os.path.exists(tokenizer_config_file):
|
||||
train_config["action_tokenizer_path"] = checkpoint_path
|
||||
|
||||
ckpt_overlay = _load_ckpt_config_overlay(checkpoint_path)
|
||||
_merge_model_fields(train_config, ckpt_overlay, overwrite=False)
|
||||
if isinstance(train_config.get("model"), dict):
|
||||
_merge_model_fields(train_config["model"], ckpt_overlay, overwrite=False)
|
||||
|
||||
return train_config
|
||||
|
||||
|
||||
def normalize_train_config_for_inference(
|
||||
train_config: dict,
|
||||
train_config_path: str,
|
||||
) -> dict:
|
||||
"""Flatten typed TrainConfig YAML into the legacy dict inference expects.
|
||||
|
||||
Serving and ``update_model_config`` still read top-level ``dof_config`` and
|
||||
nested ``data.*`` fields. New training yamls store those under ``task:``.
|
||||
"""
|
||||
try:
|
||||
from wall_x.config.loader import load_config
|
||||
|
||||
typed_cfg = load_config(train_config_path)
|
||||
except Exception as e:
|
||||
logger.debug(
|
||||
"Train config is not typed TrainConfig schema (%s): %s",
|
||||
train_config_path,
|
||||
e,
|
||||
)
|
||||
_mirror_task_fields(train_config, train_config.get("task") or {})
|
||||
_canonicalize_task_layouts(train_config)
|
||||
return train_config
|
||||
|
||||
import dataclasses
|
||||
|
||||
normalized = typed_cfg.build_data_loader_dict()
|
||||
model_dict = dataclasses.asdict(typed_cfg.model)
|
||||
_merge_model_fields(normalized, model_dict, overwrite=True)
|
||||
|
||||
normalized["data"] = dict(normalized.get("data") or {})
|
||||
if isinstance(train_config.get("model"), dict):
|
||||
_merge_model_fields(normalized, train_config["model"], overwrite=False)
|
||||
for key, value in train_config.items():
|
||||
if key in _INFERENCE_MODEL_KEYS or key == "qwen_vl_act_config_path":
|
||||
if value is not None:
|
||||
normalized[key] = value
|
||||
elif key == "data" and isinstance(value, dict):
|
||||
normalized["data"].update(value)
|
||||
|
||||
_mirror_task_fields(normalized, dataclasses.asdict(typed_cfg.task))
|
||||
_canonicalize_task_layouts(normalized)
|
||||
return normalized
|
||||
|
||||
|
||||
def register_data_backend(train_config: dict) -> None:
|
||||
"""Register the data backend before typed data config construction."""
|
||||
from wall_x.data._registry import _set_data_backend
|
||||
|
||||
data_section = train_config.get("data", {})
|
||||
dataset_type = train_config.get("dataset_type") or data_section.get(
|
||||
"dataset_type", "lerobot"
|
||||
)
|
||||
_set_data_backend(dataset_type)
|
||||
|
||||
|
||||
def build_model_config(
|
||||
config_class,
|
||||
checkpoint_path: str,
|
||||
train_config: dict,
|
||||
train_config_path: str | None = None,
|
||||
):
|
||||
"""Build the HF model config and inject train_config-derived fields.
|
||||
|
||||
``config_class`` is supplied by the variant adapter. If the checkpoint does
|
||||
not include ``config.json``, this falls back to
|
||||
``train_config["qwen_vl_act_config_path"]``.
|
||||
"""
|
||||
ckpt_config_path = os.path.join(checkpoint_path, "config.json")
|
||||
if os.path.exists(ckpt_config_path):
|
||||
resolved = ckpt_config_path
|
||||
else:
|
||||
resolved = train_config.get("qwen_vl_act_config_path")
|
||||
if resolved is None or not os.path.exists(resolved):
|
||||
raise ValueError(
|
||||
f"cannot load model config: checkpoint file {ckpt_config_path} "
|
||||
f"does not exist, and fallback qwen_vl_act_config_path={resolved!r} "
|
||||
"does not exist either"
|
||||
)
|
||||
train_config["qwen_vl_act_config_path"] = resolved
|
||||
|
||||
if resolved.endswith(".json"):
|
||||
model_config = config_class.from_json_file(resolved)
|
||||
else:
|
||||
model_config = config_class.from_pretrained(resolved)
|
||||
|
||||
legacy_train_config = (
|
||||
normalize_train_config_for_inference(train_config, train_config_path)
|
||||
if train_config_path
|
||||
else train_config
|
||||
)
|
||||
model_config.update_model_config(legacy_train_config)
|
||||
model_config._attn_implementation = "sdpa"
|
||||
model_config.vision_config._attn_implementation = "flash_attention_2"
|
||||
|
||||
return model_config
|
||||
|
||||
|
||||
def build_data_config(train_config_path: str, train_config: dict):
|
||||
"""Build the data config used by image resizing and preprocessing.
|
||||
|
||||
New checkpoints use the typed schema. Older checkpoints may still use a flat
|
||||
legacy schema, so this falls back to the active data backend.
|
||||
"""
|
||||
from wall_x.config.loader import load_config
|
||||
from wall_x.data import data_backend
|
||||
|
||||
try:
|
||||
typed_cfg = load_config(train_config_path)
|
||||
backend = data_backend()
|
||||
if backend.supports("load_trainer_data_config"):
|
||||
return backend.load_trainer_data_config(typed_cfg)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"typed TrainConfig loading failed; falling back to raw data config: %s",
|
||||
e,
|
||||
)
|
||||
|
||||
backend = data_backend()
|
||||
if backend.supports("load_trainer_data_config_from_yaml_dict"):
|
||||
return backend.load_trainer_data_config_from_yaml_dict(train_config)
|
||||
raise RuntimeError(
|
||||
f"active data backend {backend!r} cannot build a trainer data config"
|
||||
)
|
||||
Reference in New Issue
Block a user