Upload folder using huggingface_hub

src/dataset.py

@@ -0,0 +1,194 @@
+"""Dataset utilities for cached A2C2 BEHAVIOR/OpenPI parquet exports."""
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+import random
+from typing import Iterator, NamedTuple
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+import torch
+from torch import Tensor
+from torch.utils.data import IterableDataset, get_worker_info
+
+
+class EpisodePair(NamedTuple):
+    data_path: Path
+    latent_path: Path
+
+
+def resolve_dataset_root(path: Path) -> Path:
+    """Resolve either an A2C2 root or its parent directory."""
+
+    path = path.expanduser().resolve()
+    if (path / "data").is_dir() and (path / "latent" / "data").is_dir():
+        return path
+
+    candidates = [p for p in path.iterdir() if (p / "data").is_dir() and (p / "latent" / "data").is_dir()]
+    if len(candidates) == 1:
+        return candidates[0].resolve()
+    if not candidates:
+        raise FileNotFoundError(f"No A2C2 dataset root found under {path}")
+    names = ", ".join(str(p) for p in candidates)
+    raise ValueError(f"Multiple dataset roots found under {path}; pass one explicitly: {names}")
+
+
+def discover_episode_pairs(dataset_root: Path, task_dir: str | None = None) -> list[EpisodePair]:
+    """Find matching data/latent parquet pairs."""
+
+    dataset_root = resolve_dataset_root(dataset_root)
+    pattern = f"{task_dir}/episode_*.parquet" if task_dir else "task-*/episode_*.parquet"
+    data_paths = sorted((dataset_root / "data").glob(pattern))
+    pairs: list[EpisodePair] = []
+    for data_path in data_paths:
+        rel = data_path.relative_to(dataset_root / "data")
+        latent_path = dataset_root / "latent" / "data" / rel
+        if not latent_path.is_file():
+            raise FileNotFoundError(f"Missing latent parquet for {data_path}: {latent_path}")
+        pairs.append(EpisodePair(data_path=data_path, latent_path=latent_path))
+    if not pairs:
+        raise FileNotFoundError(f"No episode parquet files found in {dataset_root / 'data'}")
+    return pairs
+
+
+def split_episode_pairs(
+    pairs: list[EpisodePair],
+    val_ratio: float,
+    seed: int,
+    max_episodes: int | None = None,
+) -> tuple[list[EpisodePair], list[EpisodePair]]:
+    """Shuffle episode pairs and split into train/validation subsets."""
+
+    pairs = list(pairs)
+    rng = random.Random(seed)
+    rng.shuffle(pairs)
+    if max_episodes is not None:
+        pairs = pairs[:max_episodes]
+    val_count = int(round(len(pairs) * val_ratio))
+    if val_ratio > 0 and val_count == 0 and len(pairs) > 1:
+        val_count = 1
+    val_pairs = pairs[:val_count]
+    train_pairs = pairs[val_count:]
+    if not train_pairs:
+        raise ValueError("No training episodes left after split.")
+    return train_pairs, val_pairs
+
+
+def fixed_or_variable_list_to_numpy(column: pa.ChunkedArray, dtype: np.dtype) -> np.ndarray:
+    """Convert Arrow list/fixed-size-list columns to dense numpy arrays."""
+
+    array = column.combine_chunks()
+    if pa.types.is_fixed_size_list(array.type):
+        outer_size = array.type.list_size
+        inner = array.values
+        if pa.types.is_fixed_size_list(inner.type):
+            inner_size = inner.type.list_size
+            flat = inner.values.to_numpy(zero_copy_only=False)
+            return np.asarray(flat, dtype=dtype).reshape(len(array), outer_size, inner_size)
+        flat = inner.to_numpy(zero_copy_only=False)
+        return np.asarray(flat, dtype=dtype).reshape(len(array), outer_size)
+    return np.asarray(array.to_pylist(), dtype=dtype)
+
+
+def load_episode(pair: EpisodePair) -> dict[str, np.ndarray]:
+    """Load one episode's state/action/chunk rows plus aligned base-policy latents."""
+
+    data = pq.read_table(
+        pair.data_path,
+        columns=[
+            "observation.state",
+            "action",
+            "a2c2.base_action_chunk",
+            "a2c2.valid_action_mask",
+        ],
+    )
+    latent = pq.read_table(pair.latent_path, columns=["a2c2.base_policy_z"])
+    if data.num_rows != latent.num_rows:
+        raise ValueError(f"Row mismatch: {pair.data_path} has {data.num_rows}, {pair.latent_path} has {latent.num_rows}")
+
+    return {
+        "states": fixed_or_variable_list_to_numpy(data.column("observation.state"), np.float32),
+        "actions": fixed_or_variable_list_to_numpy(data.column("action"), np.float32),
+        "chunks": fixed_or_variable_list_to_numpy(data.column("a2c2.base_action_chunk"), np.float32),
+        "masks": fixed_or_variable_list_to_numpy(data.column("a2c2.valid_action_mask"), np.bool_),
+        "zs": fixed_or_variable_list_to_numpy(latent.column("a2c2.base_policy_z"), np.float32),
+    }
+
+
+class A2C2RandomSampleDataset(IterableDataset):
+    """Randomly sample valid (source frame t, chunk offset k) training examples."""
+
+    def __init__(
+        self,
+        episode_pairs: list[EpisodePair],
+        action_horizon: int,
+        samples_per_episode: int,
+        seed: int,
+        total_samples: int | None = None,
+    ) -> None:
+        super().__init__()
+        self.episode_pairs = list(episode_pairs)
+        self.action_horizon = action_horizon
+        self.samples_per_episode = samples_per_episode
+        self.seed = seed
+        self.total_samples = total_samples
+
+    def __iter__(self) -> Iterator[dict[str, np.ndarray]]:
+        worker = get_worker_info()
+        worker_id = worker.id if worker else 0
+        num_workers = worker.num_workers if worker else 1
+        pairs = self.episode_pairs[worker_id::num_workers]
+        if not pairs:
+            return
+
+        rng = np.random.default_rng(self.seed + worker_id)
+        yielded = 0
+        while self.total_samples is None or yielded < self.total_samples:
+            order = rng.permutation(len(pairs))
+            for episode_idx in order:
+                episode = load_episode(pairs[int(episode_idx)])
+                rows = episode["actions"].shape[0]
+                for _ in range(self.samples_per_episode):
+                    if self.total_samples is not None and yielded >= self.total_samples:
+                        return
+                    source_idx = int(rng.integers(0, rows))
+                    valid_offsets = np.flatnonzero(episode["masks"][source_idx])
+                    if valid_offsets.size == 0:
+                        continue
+                    k = int(rng.choice(valid_offsets))
+                    target_idx = source_idx + k
+                    if target_idx >= rows:
+                        continue
+
+                    base_action = episode["chunks"][source_idx, k]
+                    expert_action = episode["actions"][target_idx]
+                    denom = max(self.action_horizon - 1, 1)
+                    phase = 2.0 * math.pi * float(k) / denom
+                    yield {
+                        "observation_state": episode["states"][target_idx],
+                        "base_action_chunk": episode["chunks"][source_idx],
+                        "base_policy_z": episode["zs"][source_idx],
+                        "time_feature": np.asarray([math.sin(phase), math.cos(phase)], dtype=np.float32),
+                        "valid_action_mask": episode["masks"][source_idx],
+                        "base_action": base_action,
+                        "target_delta": expert_action - base_action,
+                        "expert_action": expert_action,
+                    }
+                    yielded += 1
+
+
+def move_batch_to_device(batch: dict[str, Tensor], device: torch.device) -> dict[str, Tensor]:
+    return {key: value.to(device, non_blocking=True) if torch.is_tensor(value) else value for key, value in batch.items()}
+
+
+def pick_device(raw: str) -> torch.device:
+    if raw != "auto":
+        return torch.device(raw)
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")

src/model.py

@@ -0,0 +1,209 @@
+"""A2C2 correction head architecture for cached BEHAVIOR/OpenPI features."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import math
+
+import torch
+from torch import Tensor, nn
+
+
+@dataclass(frozen=True)
+class A2C2CorrectionHeadConfig:
+    state_dim: int = 256
+    action_dim: int = 23
+    action_horizon: int = 32
+    base_policy_z_dim: int = 2048
+    use_base_policy_z: bool = True
+    time_dim: int = 2
+
+    dim_model: int = 512
+    n_heads: int = 8
+    n_encoder_layers: int = 6
+    dim_feedforward: int = 2048
+    dropout: float = 0.1
+    mlp_hidden_dim: int = 1024
+
+
+def _sinusoidal_positions(length: int, dim: int) -> Tensor:
+    if dim % 2 != 0:
+        raise ValueError("dim must be even for sinusoidal positional encoding.")
+
+    position = torch.arange(length, dtype=torch.float32).unsqueeze(1)
+    div_term = torch.exp(torch.arange(0, dim, 2, dtype=torch.float32) * (-math.log(10000.0) / dim))
+    pe = torch.zeros(length, dim, dtype=torch.float32)
+    pe[:, 0::2] = torch.sin(position * div_term)
+    pe[:, 1::2] = torch.cos(position * div_term)
+    return pe
+
+
+class A2C2CorrectionHead(nn.Module):
+    """Transformer + MLP correction head following the A2C2 residual design."""
+
+    def __init__(self, config: A2C2CorrectionHeadConfig | None = None) -> None:
+        super().__init__()
+        self.config = config or A2C2CorrectionHeadConfig()
+        cfg = self.config
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, cfg.dim_model))
+        self.type_embedding = nn.Parameter(torch.zeros(6, cfg.dim_model))
+
+        self.state_proj = nn.Linear(cfg.state_dim, cfg.dim_model)
+        if cfg.use_base_policy_z:
+            self.z_proj = nn.Linear(cfg.base_policy_z_dim, cfg.dim_model)
+        self.time_proj = nn.Linear(cfg.time_dim, cfg.dim_model)
+        self.action_proj = nn.Linear(cfg.action_dim, cfg.dim_model)
+
+        chunk_pos = _sinusoidal_positions(cfg.action_horizon, cfg.dim_model)
+        self.register_buffer("chunk_pos_embedding", chunk_pos, persistent=False)
+
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=cfg.dim_model,
+            nhead=cfg.n_heads,
+            dim_feedforward=cfg.dim_feedforward,
+            dropout=cfg.dropout,
+            activation="gelu",
+            batch_first=True,
+            norm_first=True,
+        )
+        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=cfg.n_encoder_layers)
+        self.encoder_norm = nn.LayerNorm(cfg.dim_model)
+
+        head_input_token_count = 5 if cfg.use_base_policy_z else 4
+        head_input_dim = cfg.dim_model * head_input_token_count + cfg.action_dim
+        self.residual_head = nn.Sequential(
+            nn.Linear(head_input_dim, cfg.mlp_hidden_dim),
+            nn.GELU(),
+            nn.Dropout(cfg.dropout),
+            nn.Linear(cfg.mlp_hidden_dim, cfg.mlp_hidden_dim),
+            nn.GELU(),
+            nn.Dropout(cfg.dropout),
+            nn.Linear(cfg.mlp_hidden_dim, cfg.action_dim),
+        )
+
+        self._reset_parameters()
+
+    @staticmethod
+    def make_time_feature(chunk_index: Tensor, horizon: int) -> Tensor:
+        """Create [sin, cos] phase features from chunk indices."""
+
+        idx = chunk_index.to(dtype=torch.float32)
+        denom = max(horizon - 1, 1)
+        phase = 2.0 * math.pi * idx / denom
+        return torch.stack([torch.sin(phase), torch.cos(phase)], dim=-1)
+
+    def forward(
+        self,
+        observation_state: Tensor,
+        selected_base_action: Tensor,
+        base_action_chunk: Tensor,
+        base_policy_z: Tensor,
+        time_feature: Tensor,
+        valid_action_mask: Tensor | None = None,
+    ) -> Tensor:
+        """Predict residual action delta.
+
+        Args:
+            observation_state: [B, state_dim]
+            selected_base_action: [B, action_dim], the current action being corrected.
+            base_action_chunk: [B, H, action_dim]
+            base_policy_z: [B, z_dim]
+            time_feature: [B, 2]
+            valid_action_mask: optional bool tensor [B, H], True for valid chunk
+                entries. Invalid chunk entries are ignored by transformer attention.
+
+        Returns:
+            Tensor [B, action_dim], the predicted residual delta.
+        """
+
+        cfg = self.config
+        batch_size = observation_state.shape[0]
+        device = observation_state.device
+        dtype = observation_state.dtype
+
+        self._validate_inputs(observation_state, selected_base_action, base_action_chunk, base_policy_z, time_feature)
+
+        cls = self.cls_token.to(device=device, dtype=dtype).expand(batch_size, -1, -1)
+        cls = cls + self.type_embedding[0].to(device=device, dtype=dtype)
+
+        state_token = self.state_proj(observation_state).unsqueeze(1)
+        state_token = state_token + self.type_embedding[1].to(device=device, dtype=dtype)
+
+        time_token = self.time_proj(time_feature).unsqueeze(1)
+        time_token = time_token + self.type_embedding[3].to(device=device, dtype=dtype)
+
+        selected_action_token = self.action_proj(selected_base_action).unsqueeze(1)
+        selected_action_token = selected_action_token + self.type_embedding[4].to(device=device, dtype=dtype)
+
+        chunk_tokens = self.action_proj(base_action_chunk)
+        chunk_pos = self.chunk_pos_embedding[: base_action_chunk.shape[1]].to(device=device, dtype=dtype)
+        chunk_tokens = chunk_tokens + chunk_pos.unsqueeze(0)
+        chunk_tokens = chunk_tokens + self.type_embedding[5].to(device=device, dtype=dtype)
+
+        prefix_tokens = [cls, state_token]
+        if cfg.use_base_policy_z:
+            z_token = self.z_proj(base_policy_z).unsqueeze(1)
+            z_token = z_token + self.type_embedding[2].to(device=device, dtype=dtype)
+            prefix_tokens.append(z_token)
+        prefix_tokens.extend([time_token, selected_action_token])
+
+        tokens = torch.cat([*prefix_tokens, chunk_tokens], dim=1)
+
+        padding_mask = None
+        if valid_action_mask is not None:
+            valid_action_mask = valid_action_mask.to(device=device, dtype=torch.bool)
+            prefix_mask = torch.zeros(batch_size, len(prefix_tokens), device=device, dtype=torch.bool)
+            padding_mask = torch.cat([prefix_mask, ~valid_action_mask], dim=1)
+
+        encoded = self.encoder(tokens, src_key_padding_mask=padding_mask)
+        encoded = self.encoder_norm(encoded)
+
+        cls_state = encoded[:, 0]
+        state_state = encoded[:, 1]
+        if cfg.use_base_policy_z:
+            z_state = encoded[:, 2]
+            time_state = encoded[:, 3]
+            selected_action_state = encoded[:, 4]
+            head_states = [cls_state, state_state, z_state, time_state, selected_action_state]
+        else:
+            time_state = encoded[:, 2]
+            selected_action_state = encoded[:, 3]
+            head_states = [cls_state, state_state, time_state, selected_action_state]
+
+        head_input = torch.cat(
+            [*head_states, selected_base_action],
+            dim=-1,
+        )
+        return self.residual_head(head_input)
+
+    def _validate_inputs(
+        self,
+        observation_state: Tensor,
+        selected_base_action: Tensor,
+        base_action_chunk: Tensor,
+        base_policy_z: Tensor,
+        time_feature: Tensor,
+    ) -> None:
+        cfg = self.config
+        if observation_state.ndim != 2 or observation_state.shape[-1] != cfg.state_dim:
+            raise ValueError(f"observation_state must have shape [B, {cfg.state_dim}].")
+        if selected_base_action.ndim != 2 or selected_base_action.shape[-1] != cfg.action_dim:
+            raise ValueError(f"selected_base_action must have shape [B, {cfg.action_dim}].")
+        if base_action_chunk.ndim != 3 or base_action_chunk.shape[-1] != cfg.action_dim:
+            raise ValueError(f"base_action_chunk must have shape [B, H, {cfg.action_dim}].")
+        if base_action_chunk.shape[1] > cfg.action_horizon:
+            raise ValueError(f"base_action_chunk horizon cannot exceed {cfg.action_horizon}.")
+        if cfg.use_base_policy_z and (base_policy_z.ndim != 2 or base_policy_z.shape[-1] != cfg.base_policy_z_dim):
+            raise ValueError(f"base_policy_z must have shape [B, {cfg.base_policy_z_dim}].")
+        if time_feature.ndim != 2 or time_feature.shape[-1] != cfg.time_dim:
+            raise ValueError(f"time_feature must have shape [B, {cfg.time_dim}].")
+
+    def _reset_parameters(self) -> None:
+        nn.init.trunc_normal_(self.cls_token, std=0.02)
+        nn.init.trunc_normal_(self.type_embedding, std=0.02)
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)