Spaces:
Running on Zero
Running on Zero
| """Componentwise Stage2 action statistics, portable artifacts, and initialization.""" | |
| from __future__ import annotations | |
| from contextlib import ExitStack | |
| import json | |
| import logging | |
| import math | |
| from pathlib import Path | |
| from tempfile import TemporaryDirectory | |
| import time | |
| import numpy as np | |
| import torch | |
| from tcp4d.action import future_actions_in_current_camera | |
| LOGGER = logging.getLogger(__name__) | |
| NORM_FILENAME = "action_norm.json" | |
| NORM_METHOD = "componentwise" | |
| NORM_VERSION = 2 | |
| POSITION_FALLBACK_SPAN = 1e-4 | |
| POSITION_MIN_SPAN = 1e-3 | |
| STAT_NAMES = ("mean", "std", "q01", "q99") | |
| ACTION_COMPONENTS = ["x", "y", "z", "rotation6d_0", "rotation6d_1", "rotation6d_2", | |
| "rotation6d_3", "rotation6d_4", "rotation6d_5", "gripper"] | |
| class ActionNormalizationError(ValueError): | |
| """Invalid action shapes, normalization parameters, or checkpoint metadata.""" | |
| def configure_action_normalization(config): | |
| """New runs have exactly one action normalization format.""" | |
| config.setdefault("action_normalization", NORM_METHOD) | |
| config.setdefault("action_normalization_version", NORM_VERSION) | |
| # Fitting the complete dataset on rank zero can exceed the default process | |
| # group timeout. This also covers index/profile construction before fitting. | |
| config.setdefault("distributed_timeout_seconds", 86400) | |
| timeout = config["distributed_timeout_seconds"] | |
| if not isinstance(timeout, (int, float)) or not math.isfinite(timeout) or timeout <= 0: | |
| raise ValueError("distributed_timeout_seconds must be finite and positive") | |
| require_action_normalization_config(config) | |
| def require_action_normalization_config(config): | |
| pair = (config.get("action_normalization"), config.get("action_normalization_version")) | |
| if pair != (NORM_METHOD, NORM_VERSION): | |
| raise ActionNormalizationError( | |
| "Stage2 requires componentwise action normalization v2; " | |
| "start new componentwise runs from Stage1." | |
| ) | |
| def componentwise_parameters(q01, q99, minimum, maximum): | |
| """XYZ affine parameters; rotation columns and gripper retain their units.""" | |
| q01, q99, minimum, maximum = [np.asarray(v, dtype=np.float64) for v in (q01, q99, minimum, maximum)] | |
| if (q01.ndim != 2 or q01.shape[0] != 1 or q01.shape[1] not in (10, 20) | |
| or any(v.shape != q01.shape or not np.isfinite(v).all() for v in (q01, q99, minimum, maximum)) | |
| or (minimum > q01).any() or (q01 > q99).any() or (q99 > maximum).any()): | |
| raise ActionNormalizationError("Invalid componentwise action statistics") | |
| fallback = q99 - q01 < POSITION_FALLBACK_SPAN | |
| low, high = np.where(fallback, minimum, q01), np.where(fallback, maximum, q99) | |
| # Each arm contributes xyz + rotation6D + gripper. Fit affine scaling only | |
| # to XYZ; rotation directions and the [-1,1] gripper score retain their units. | |
| position = np.arange(q01.shape[1]) % 10 < 3 | |
| offset, scale = np.zeros_like(q01), np.ones_like(q01) | |
| offset[:, position] = (low[:, position] + high[:, position]) / 2 | |
| scale[:, position] = np.maximum(high[:, position] - low[:, position], POSITION_MIN_SPAN) / 2 | |
| return offset, scale | |
| def _affine_parameters(action, offset, scale): | |
| if torch.is_tensor(action): | |
| if not action.is_floating_point(): | |
| raise ActionNormalizationError("Actions must use a floating dtype") | |
| action = action.float() | |
| offset = torch.as_tensor(offset, dtype=torch.float32, device=action.device) | |
| scale = torch.as_tensor(scale, dtype=torch.float32, device=action.device) | |
| finite = torch.isfinite | |
| else: | |
| action = np.asarray(action) | |
| if not np.issubdtype(action.dtype, np.floating): | |
| raise ActionNormalizationError("Actions must use a floating dtype") | |
| offset, scale = np.asarray(offset, dtype=np.float32), np.asarray(scale, dtype=np.float32) | |
| finite = np.isfinite | |
| if (action.ndim < 2 or offset.shape != (1, action.shape[-1]) or scale.shape != offset.shape | |
| or not bool(finite(offset).all() & finite(scale).all() & (scale > 0).all())): | |
| raise ActionNormalizationError("Expected finite affine [1,D] parameters with positive scales") | |
| return action, offset, scale | |
| def componentwise_normalize(action, offset, scale): | |
| action, offset, scale = _affine_parameters(action, offset, scale) | |
| return (action - offset) / scale | |
| def componentwise_unnormalize(action, offset, scale): | |
| action, offset, scale = _affine_parameters(action, offset, scale) | |
| return action * scale + offset | |
| def action_layout(num_arms): | |
| arms = ["left", "right"] if num_arms == 2 else [f"arm_{i}" for i in range(num_arms)] | |
| return [f"{arm}.{component}" for arm in arms for component in ACTION_COMPONENTS] | |
| def iter_future_action_rows(states, extrinsics, anchors, segment_ends, horizon, *, chunk_size=8192): | |
| """Yield every real (anchor, future offset) label as float32 [N,A*10]. | |
| This uses the same coordinate/rotation/gripper conversion as sampled labels. | |
| Extrinsics may be a static camera matrix or a per-frame camera trajectory. | |
| No padded rows or image reads are involved. | |
| """ | |
| anchors = np.asarray(anchors, dtype=np.int64) | |
| segment_ends = np.asarray(segment_ends, dtype=np.int64) | |
| if anchors.shape != segment_ends.shape or horizon < 1 or chunk_size < 1: | |
| raise ValueError("Invalid action statistics window specification") | |
| static_camera = np.ndim(extrinsics) == 2 | |
| for offset in range(1, horizon + 1): | |
| selected = anchors[anchors + offset < segment_ends] | |
| for begin in range(0, len(selected), chunk_size): | |
| current_indices = selected[begin:begin + chunk_size] | |
| future_indices = current_indices + offset | |
| state = torch.tensor(np.asarray(states[future_indices]), dtype=torch.float32) | |
| if static_camera: | |
| current = torch.tensor(np.asarray(extrinsics), dtype=torch.float32) | |
| future = current.expand(len(future_indices), -1, -1) | |
| else: | |
| current = torch.tensor(np.asarray(extrinsics[current_indices]), dtype=torch.float32) | |
| future = torch.tensor(np.asarray(extrinsics[future_indices]), dtype=torch.float32) | |
| yield future_actions_in_current_camera(state, future, current).flatten(-2).numpy() | |
| def validate_action_norm(record, num_arms): | |
| require_action_normalization_config({"action_normalization": record.get("method"), | |
| "action_normalization_version": record.get("version")}) | |
| expected = {"horizon_shared": True, "num_arms": num_arms, "action_dim": num_arms * 10, | |
| "layout": action_layout(num_arms), "statistics_scope": "train_valid_window_actions"} | |
| expected.update(position_fallback_span=POSITION_FALLBACK_SPAN, position_min_span=POSITION_MIN_SPAN, | |
| component_rules=["quantile_xyz", "identity_rotation6d", "identity_gripper"]) | |
| for key, value in expected.items(): | |
| if record.get(key) != value: | |
| raise ActionNormalizationError(f"Invalid action normalization {key}: expected {value!r}") | |
| if not isinstance(record.get("count"), int) or record["count"] < 1: | |
| raise ActionNormalizationError("Action normalization requires a positive sample count") | |
| stats = {} | |
| for name in (*STAT_NAMES, "min", "max", "offset", "scale"): | |
| try: | |
| value = np.asarray(record[name], dtype=np.float64) | |
| except (KeyError, TypeError, ValueError) as exc: | |
| raise ActionNormalizationError(f"Invalid action normalization {name}") from exc | |
| if value.shape != (1, num_arms * 10) or not np.isfinite(value).all(): | |
| raise ActionNormalizationError(f"{name} must contain finite [1,{num_arms * 10}] statistics") | |
| stats[name] = value | |
| if (stats["std"] < 0).any() or (stats["q99"] < stats["q01"]).any(): | |
| raise ActionNormalizationError("Invalid standard deviation or reversed quantiles") | |
| offset, scale = componentwise_parameters(stats["q01"], stats["q99"], stats["min"], stats["max"]) | |
| for name, value in (("offset", offset), ("scale", scale)): | |
| if not np.array_equal(stats[name].astype(np.float32), value.astype(np.float32)): | |
| raise ActionNormalizationError(f"Saved {name} disagrees with componentwise rules") | |
| scale = scale.astype(np.float32) | |
| if not np.isfinite(offset.astype(np.float32)).all(): | |
| raise ActionNormalizationError("Float32 offset must be finite") | |
| if not np.isfinite(scale).all() or not (scale > 0).all(): | |
| raise ActionNormalizationError("Float32 normalization scale must be finite and positive") | |
| return stats | |
| def compute_action_norm(mixture, num_arms, scratch_dir): | |
| """Exact XYZ-only quantiles with bounded RAM and one temporary file per XYZ axis. | |
| Every valid window action has equal weight, regardless of sampler weights. | |
| In-place partition on contiguous memory-mapped columns avoids a full RAM copy. | |
| """ | |
| dimension = num_arms * 10 | |
| position_indices = np.flatnonzero(np.arange(dimension) % 10 < 3) | |
| statistics_dim = len(position_indices) | |
| layout = action_layout(num_arms) | |
| minimum, maximum = np.full(statistics_dim, np.inf), np.full(statistics_dim, -np.inf) | |
| count = 0 | |
| mean, m2 = np.zeros(statistics_dim), np.zeros(statistics_dim) | |
| sources = [] | |
| last_progress = time.monotonic() | |
| with TemporaryDirectory(prefix=".action-norm-", dir=scratch_dir) as temporary: | |
| paths = [Path(temporary) / f"xyz-{i}.bin" for i in position_indices] | |
| with ExitStack() as stack: | |
| handles = [stack.enter_context(path.open("wb")) for path in paths] | |
| for source in mixture.sources: | |
| if source.weight <= 0: | |
| continue | |
| dataset = source.dataset | |
| if dataset.split != "train": | |
| raise ActionNormalizationError("Quantile fitting may only use the training split") | |
| source_count = 0 | |
| LOGGER.info("Fit XYZ normalization: source=%s, streams=%d, position dimensions=%d", | |
| source.name, len(dataset), statistics_dim) | |
| for values in dataset.iter_action_targets(): | |
| values = np.asarray(values, dtype=np.float32) | |
| if values.ndim != 2 or values.shape[1] != dimension or not np.isfinite(values).all(): | |
| raise ActionNormalizationError(f"Invalid action labels in source {source.name}") | |
| size = len(values) | |
| if not size: | |
| continue | |
| values = values[:, position_indices] | |
| minimum = np.minimum(minimum, values.min(axis=0)) | |
| maximum = np.maximum(maximum, values.max(axis=0)) | |
| batch = values.astype(np.float64) | |
| batch_mean = batch.mean(axis=0) | |
| delta = batch_mean - mean | |
| new_count = count + size | |
| m2 += np.square(batch - batch_mean).sum(axis=0) + delta**2 * (count * size / new_count) | |
| mean += delta * (size / new_count) | |
| count = new_count | |
| source_count += size | |
| for column, handle in enumerate(handles): | |
| np.ascontiguousarray(values[:, column]).tofile(handle) | |
| if time.monotonic() - last_progress >= 30: | |
| LOGGER.info("Fit XYZ normalization: %d labels, %.2f GiB temporary XYZ columns", | |
| count, count * statistics_dim * 4 / 1024**3) | |
| last_progress = time.monotonic() | |
| LOGGER.info("Fit XYZ normalization: source=%s, valid window actions=%d", source.name, source_count) | |
| sources.append({"name": source.name, "count": source_count, | |
| "history_frames": dataset.history_frames, | |
| "prediction_horizon": dataset.prediction_horizon, | |
| "selection": dataset.manifest()}) | |
| if count == 0: | |
| raise ActionNormalizationError("No valid training actions for XYZ normalization") | |
| ranks = np.array([0.01, 0.99]) * (count - 1) | |
| lower, upper = np.floor(ranks).astype(np.int64), np.ceil(ranks).astype(np.int64) | |
| quantiles = np.empty((2, statistics_dim), dtype=np.float64) | |
| for column, path in enumerate(paths): | |
| LOGGER.info("Select exact XYZ quantiles: %s (%d/%d)", | |
| layout[position_indices[column]], column + 1, statistics_dim) | |
| values = np.memmap(path, mode="r+", dtype=np.float32, shape=(count,)) | |
| values.partition(np.unique(np.r_[lower, upper])) | |
| lo, hi = values[lower].astype(np.float64), values[upper].astype(np.float64) | |
| quantiles[:, column] = lo + (hi - lo) * (ranks - lower) | |
| del values | |
| def action_columns(values): | |
| # Keep the action-aligned artifact layout. Non-XYZ slots are unused zeros, | |
| # not measured statistics; their actual affine transform stays (0, 1). | |
| expanded = np.zeros((*values.shape[:-1], dimension), dtype=np.float64) | |
| expanded[..., position_indices] = values | |
| return expanded | |
| mean, m2 = action_columns(mean), action_columns(m2) | |
| minimum, maximum = action_columns(minimum), action_columns(maximum) | |
| quantiles = action_columns(quantiles) | |
| record = { | |
| "version": NORM_VERSION, "method": NORM_METHOD, | |
| "horizon_shared": True, "num_arms": num_arms, "action_dim": dimension, | |
| "layout": action_layout(num_arms), "statistics_scope": "train_valid_window_actions", | |
| "count": count, "sources": sources, | |
| "mean": mean[None].tolist(), "std": np.sqrt(np.maximum(m2 / count, 0))[None].tolist(), | |
| "q01": quantiles[0:1].tolist(), "q99": quantiles[1:2].tolist(), | |
| } | |
| offset, scale = componentwise_parameters(quantiles[0:1], quantiles[1:2], minimum[None], maximum[None]) | |
| record.update(min=minimum[None].tolist(), max=maximum[None].tolist(), | |
| offset=offset.tolist(), scale=scale.tolist(), | |
| position_fallback_span=POSITION_FALLBACK_SPAN, position_min_span=POSITION_MIN_SPAN, | |
| component_rules=["quantile_xyz", "identity_rotation6d", "identity_gripper"]) | |
| validate_action_norm(record, num_arms) | |
| return record | |
| def save_action_norm(directory, record): | |
| validate_action_norm(record, record["num_arms"]) | |
| path = Path(directory) / NORM_FILENAME | |
| temporary = path.with_suffix(".json.tmp") | |
| temporary.write_text(json.dumps(record, indent=2, allow_nan=False), encoding="utf-8") | |
| temporary.replace(path) | |
| def load_checkpoint_action_norm(checkpoint, num_arms): | |
| path = Path(checkpoint).expanduser() | |
| directory = path if path.is_dir() else path.parent | |
| config_path = directory / "config.json" | |
| if not config_path.is_file(): | |
| raise ActionNormalizationError(f"Componentwise Stage2 checkpoint requires {config_path}") | |
| config = json.loads(config_path.read_text(encoding="utf-8")) | |
| require_action_normalization_config(config) | |
| path = directory / NORM_FILENAME | |
| if not path.is_file(): | |
| raise ActionNormalizationError(f"Componentwise Stage2 checkpoint requires {path}") | |
| record = json.loads(path.read_text(encoding="utf-8")) | |
| if (record.get("method"), record.get("version")) != (config["action_normalization"], config["action_normalization_version"]): | |
| raise ActionNormalizationError("Checkpoint config and action_norm.json versions disagree") | |
| validate_action_norm(record, num_arms) | |
| return record | |
| def assert_action_norm_matches(policy, record): | |
| stats = validate_action_norm(record, policy.num_arms) | |
| for name in ("offset", "scale"): | |
| actual = getattr(policy, f"action_{name}").detach().float().cpu() | |
| expected = torch.tensor(stats[name], dtype=torch.float32) | |
| if actual.shape != expected.shape or not torch.equal(actual, expected): | |
| raise ActionNormalizationError(f"Checkpoint action_{name} differs from {NORM_FILENAME}") | |
| def initialize_action_normalization(policy, dataset, config, accelerator): | |
| """Fit only on rank zero, or reuse a checkpoint; publish one identical record.""" | |
| from accelerate.utils import broadcast_object_list | |
| checkpoint = config.get("resume") | |
| result = [None] | |
| if accelerator.is_main_process: | |
| try: | |
| root = Path(config["output_dir"]) | |
| root.mkdir(parents=True, exist_ok=True) | |
| if checkpoint: | |
| record = load_checkpoint_action_norm(checkpoint, policy.num_arms) | |
| assert_action_norm_matches(policy, record) | |
| else: | |
| record = compute_action_norm(dataset, policy.num_arms, root) | |
| save_action_norm(root, record) | |
| result[0] = {"record": record} | |
| except Exception as exc: | |
| # Broadcast failures too, so other ranks do not wait for missing data. | |
| result[0] = {"error": f"{type(exc).__name__}: {exc}"} | |
| broadcast_object_list(result) | |
| if "error" in result[0]: | |
| raise ActionNormalizationError(result[0]["error"]) | |
| record = result[0]["record"] | |
| if checkpoint: | |
| assert_action_norm_matches(policy, record) | |
| else: | |
| policy.set_action_component_stats(record["offset"], record["scale"]) | |
| LOGGER.info("Action normalization: %s, %d XYZ dimensions scaled; rotation6D/gripper unchanged; " | |
| "%d valid training window actions%s", record["method"], policy.num_arms * 3, | |
| record["count"], " (reused checkpoint)" if checkpoint else "") | |
| return record | |