Got it — adding on-the-fly normalization and excavator sensor layout

I'll modify data/excavator_dataset.py to: 1) accept a sensor_layout list (names + dims), 2) compute running per-channel mean/std from dataset files (lightweight precompute) and apply z-score normalization on the fly, 3) allow an override normalization file, and 4) keep augmentations and FP16 support. Below is the updated file.
Updated data/excavator_dataset.py

import os
import glob
import numpy as np
import torch
from torch.utils.data import Dataset


# Example excavator sensor layout: list of (name, dim) tuples.
# Adjust to your dataset channels/order.
DEFAULT_SENSOR_LAYOUT = [
    ("joint_positions", 6),   # 6 actuated joints
    ("joint_velocities", 6),  # velocities
    ("base_imu_acc", 3),      # IMU linear accel
    ("base_imu_gyro", 3),     # IMU gyro
    ("force_torque", 6),      # end-effector wrench
    ("gps_xyz", 3),           # optional global position
]


def _layout_dim(layout):
    return sum(d for _, d in layout)


class ExcavatorTrajectoryDataset(Dataset):
    """
    Expects trajectory files saved as NumPy .npz with keys:
      - observations: (T, obs_dim)  # sensor vector in the layout order
      - actions: (T, act_dim)

    This dataset returns fixed-length windows of length seq_len.
    Supports on-the-fly z-score normalization with running stats precomputed
    from the dataset (or loaded from a provided npz file).
    """

    def __init__(self,
                 data_dir,
                 seq_len=64,
                 stride=1,
                 augment=True,
                 dtype=torch.float32,
                 sensor_layout=None,
                 norm_stats_path=None,
                 eps=1e-6):
        self.files = sorted(glob.glob(os.path.join(data_dir, "*.npz")))
        assert self.files, f"No .npz files found in {data_dir}"
        self.seq_len = seq_len
        self.stride = stride
        self.augment = augment
        self.dtype = dtype
        self.eps = eps

        self.sensor_layout = sensor_layout if sensor_layout is not None else DEFAULT_SENSOR_LAYOUT
        self.obs_dim = _layout_dim(self.sensor_layout)

        # Precompute windows mapping
        self.index = []  # list of (file_idx, start)
        for fi, fpath in enumerate(self.files):
            with np.load(fpath) as d:
                T = d['observations'].shape[0]
            for start in range(0, max(1, T - seq_len + 1), stride):
                self.index.append((fi, start))

        # Load or compute normalization stats (mean, std) per channel
        if norm_stats_path is not None and os.path.exists(norm_stats_path):
            with np.load(norm_stats_path) as s:
                self.mean = s['mean']
                self.std = s['std']
            assert self.mean.shape[0] == self.obs_dim
        else:
            # Lightweight two-pass estimate: compute mean and std across windows
            self.mean, self.std = self._compute_norm_stats(sample_windows=2000)

    def __len__(self):
        return len(self.index)

    def _compute_norm_stats(self, sample_windows=1000):
        # Sample up to sample_windows random windows across dataset to estimate mean/std
        rng = np.random.default_rng(123)
        sample_idxs = rng.choice(len(self.index), size=min(sample_windows, len(self.index)), replace=False)
        sums = np.zeros(self.obs_dim, dtype=np.float64)
        sqs = np.zeros(self.obs_dim, dtype=np.float64)
        count = 0
        for idx in sample_idxs:
            fi, start = self.index[int(idx)]
            fpath = self.files[fi]
            with np.load(fpath) as d:
                obs = d['observations'][start:start + self.seq_len].astype(np.float64)
            sums += obs.sum(axis=0)
            sqs += (obs ** 2).sum(axis=0)
            count += obs.shape[0]
        mean = (sums / count).astype(np.float32)
        var = (sqs / count) - (mean.astype(np.float64) ** 2)
        std = np.sqrt(np.maximum(var, 1e-6)).astype(np.float32)
        return mean, std

    def _load_window(self, fpath, start):
        with np.load(fpath) as d:
            obs = d['observations']
            acts = d['actions']
        window_obs = obs[start:start + self.seq_len]
        window_acts = acts[start:start + self.seq_len]
        return window_obs, window_acts

    def _augment(self, obs):
        # Augmentations appropriate for sensor vectors
        aug = obs.copy()
        # gaussian noise scaled per-channel
        noise_scale = 0.005 * (np.std(obs, axis=0) + 1e-6)  # smaller noise for normalized data
        aug += np.random.randn(*aug.shape) * noise_scale
        # channel dropout
        if aug.shape[-1] > 4:
            n_ch = aug.shape[-1]
            drop_p = 0.03
            mask = (np.random.rand(n_ch) > drop_p).astype(aug.dtype)
            aug = aug * mask[None, :]
        return aug

    def _normalize(self, obs):
        # obs: (S, obs_dim)
        return (obs - self.mean[None, :]) / (self.std[None, :] + self.eps)

    def __getitem__(self, idx):
        fi, start = self.index[idx]
        fpath = self.files[fi]
        obs, acts = self._load_window(fpath, start)

        # Validate obs dimension
        if obs.shape[1] != self.obs_dim:
            raise ValueError(f"Observation dim mismatch: expected {self.obs_dim}, got {obs.shape[1]} in {fpath}")

        # Primary view (anchor) and normalization
        obs_anchor = obs.astype(np.float32)
        obs_anchor = self._normalize(obs_anchor)

        # Positive/augmented view
        if self.augment:
            obs_pos = self._augment(obs)
            obs_pos = self._normalize(obs_pos.astype(np.float32))
        else:
            obs_pos = obs_anchor.copy()

        obs_anchor_t = torch.as_tensor(obs_anchor, dtype=self.dtype)
        obs_pos_t = torch.as_tensor(obs_pos, dtype=self.dtype)
        acts_t = torch.as_tensor(acts.astype(np.float32), dtype=self.dtype)

        return {
            'obs_anchor': obs_anchor_t,   # (S, obs_dim)
            'obs_pos': obs_pos_t,         # (S, obs_dim)
            'actions': acts_t             # (S, act_dim)
        }


# Collate unchanged
def excavator_collate(batch):
    batch_anchor = torch.stack([b['obs_anchor'] for b in batch], dim=0)
    batch_pos = torch.stack([b['obs_pos'] for b in batch], dim=0)
    batch_actions = torch.stack([b['actions'] for b in batch], dim=0)
    return {
        'obs_anchor': batch_anchor,  # (B, S, obs_dim)
        'obs_pos': batch_pos,
        'actions': batch_actions
    }


if __name__ == '__main__':
    # Smoke test
    data_dir = '/path/to/mujoco_excavator_trajs'
    ds = ExcavatorTrajectoryDataset(data_dir, seq_len=32, augment=True)
    from torch.utils.data import DataLoader
    dl = DataLoader(ds, batch_size=4, collate_fn=excavator_collate)
    batch = next(iter(dl))
    print('obs_anchor', batch['obs_anchor'].shape)
    print('obs_pos', batch['obs_pos'].shape)
    print('actions', batch['actions'].shape)
Would you like me to also include a helper script to precompute and save normalization stats to disk (norm_stats.npz) so you can reuse them for deployment?
Learning: Implementing machine learning for robotics. World models and behaviors
