feat: AR rollout loss for ForwardDynamicsPredictionObjective

.typos.toml

@@ -0,0 +1,5 @@
+[default.extend-words]
+# numpy / PyTorch array-range function — not a misspelling of "arrange"
+arange = "arange"
+# "nd" abbreviation (n-dimensional) — not a misspelling of "and"
+nd = "nd"

config/model/yaak/control_transformer/raw_vjepa.yaml

@@ -379,6 +379,13 @@ objectives:
       # null is intentional: the speed projection LayerNorms its inputs and
       # foresight features pass through their own projection
       norm: null
+      # AR rollout loss: set ar_steps: 0 (or feedback_projection: null) to disable.
+      # When disabled the loss dict is unchanged — no structural diff vs the original.
+      ar_steps: 1
+      feedback_projection:
+        _target_: rmind.components.nn.Linear
+        in_features: ${image_embedding_dim}
+        out_features: ${encoder_embedding_dim}
       patch_pos_embed:
         _target_: rmind.components.position_encoding.PatchPositionEmbedding2D
         grid_size: [16, 16]

notebooks/ar_sanity_checks.ipynb

@@ -0,0 +1,341 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "7fb27b941602401d91542211134fc71a",
+   "metadata": {},
+   "source": [
+    "# AR Rollout Loss — Sanity Checks\n",
+    "\n",
+    "Two checks:\n",
+    "1. **Rollout error curve** — L1 error vs prediction step index for TF vs AR.\n",
+    "   A well-trained AR loss should show slower error growth than TF-only.\n",
+    "3. **Feature distribution overlap** — PCA of predicted vs true patch embeddings.\n",
+    "   AR predictions should cluster closer to the true distribution than TF.\n",
+    "\n",
+    "Requires a checkpoint trained with `ar_steps: 1` and a small validation batch."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "acae54e37e7d407bbb7b55eff062a284",
+   "metadata": {},
+   "outputs": [],
+   "source": "from __future__ import annotations\n\nfrom pathlib import Path\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nimport torch.nn.functional as F\nfrom einops import repeat\nfrom sklearn.decomposition import PCA\n\nfrom rmind.components.base import Modality\nfrom rmind.components.transformer.decoder import CrossAttentionDecoderHead\nfrom rmind.models.control_transformer import ControlTransformer\n\n# ── user config ────────────────────────────────────────────────────────────────\nCHECKPOINT = Path(\"../checkpoints/ar_rollout/last.ckpt\")  # adjust to your path\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nN_BATCHES = 4\n# ───────────────────────────────────────────────────────────────────────────────"
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9a63283cbaf04dbcab1f6479b197f3a8",
+   "metadata": {},
+   "source": [
+    "## Load model + datamodule"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8dd0d8092fe74a7c96281538738b07e2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "\n",
+    "from hydra import compose, initialize_config_dir\n",
+    "from hydra.utils import instantiate\n",
+    "\n",
+    "model = ControlTransformer.load_from_checkpoint(CHECKPOINT, map_location=DEVICE)\n",
+    "model.eval()\n",
+    "\n",
+    "sys.path.insert(0, \"..\")\n",
+    "with initialize_config_dir(\n",
+    "    config_dir=str(Path(\"../config\").resolve()), version_base=None\n",
+    "):\n",
+    "    cfg = compose(\n",
+    "        config_name=\"train\",\n",
+    "        overrides=[\n",
+    "            \"+experiment=yaak/control_transformer/pretrain_vjepa\",\n",
+    "            \"data.num_workers=0\",\n",
+    "            \"data.batch_size=2\",\n",
+    "        ],\n",
+    "    )\n",
+    "\n",
+    "datamodule = instantiate(cfg.data)\n",
+    "datamodule.setup(\"validate\")\n",
+    "val_loader = datamodule.val_dataloader()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "72eea5119410473aa328ad9291626812",
+   "metadata": {},
+   "source": [
+    "## Collect TF and AR predictions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8edb47106e1a46a883d545849b8ab81b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typing import TYPE_CHECKING, Any, cast\n",
+    "\n",
+    "if TYPE_CHECKING:\n",
+    "    from rmind.components.objectives.forward_dynamics import (\n",
+    "        ForwardDynamicsPredictionObjective,\n",
+    "    )\n",
+    "\n",
+    "all_tf_preds: list[torch.Tensor] = []  # [B, T-1, N, D]\n",
+    "all_ar_preds: list[torch.Tensor] = []  # [B, T-2, N, D]\n",
+    "all_targets: list[torch.Tensor] = []  # [B, T-1, N, D]\n",
+    "\n",
+    "objective = cast(\n",
+    "    \"ForwardDynamicsPredictionObjective\", model.objectives[\"forward_dynamics\"]\n",
+    ")\n",
+    "foresight_cam = (\"foresight\", \"cam_front_left\")\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    for i, raw_batch in enumerate(val_loader):\n",
+    "        if i >= N_BATCHES:\n",
+    "            break\n",
+    "\n",
+    "        batch = {\n",
+    "            k: v.to(DEVICE) if isinstance(v, torch.Tensor) else v\n",
+    "            for k, v in raw_batch.items()\n",
+    "        }\n",
+    "        episode = model.episode_builder(batch)\n",
+    "        embedding = model.encoder(\n",
+    "            src=episode.embeddings_unpacked, mask=episode.attention_mask\n",
+    "        )\n",
+    "\n",
+    "        metrics: Any = objective.compute_metrics(episode=episode, embedding=embedding)\n",
+    "\n",
+    "        tf_logits: torch.Tensor = metrics[\"_artifacts\"][\"last_embeddings\"][\n",
+    "            Modality.FORESIGHT\n",
+    "        ][\"cam_front_left\"]\n",
+    "        targets: torch.Tensor = metrics[\"_artifacts\"][\"last_targets\"][\n",
+    "            Modality.FORESIGHT\n",
+    "        ][\"cam_front_left\"]\n",
+    "\n",
+    "        all_tf_preds.append(tf_logits.cpu())\n",
+    "        all_targets.append(targets.cpu())\n",
+    "\n",
+    "        if objective.ar_steps > 0 and objective.feedback_projection is not None:\n",
+    "            index = episode.index[:-1]\n",
+    "            k = (Modality.SUMMARY, \"action_summary\")\n",
+    "            action_summary = index.select(k).parse(embedding).get(k)\n",
+    "\n",
+    "            _, _, n_patches, _ = episode.embeddings.get((\n",
+    "                Modality.IMAGE,\n",
+    "                \"cam_front_left\",\n",
+    "            )).shape\n",
+    "            mask_tokens = repeat(\n",
+    "                episode.embeddings.get((Modality.UTILITY, \"mask\"))[:, 1:],\n",
+    "                \"b t 1 d -> b t n d\",\n",
+    "                n=n_patches,\n",
+    "            )\n",
+    "            if objective.patch_pos_embed is not None:\n",
+    "                mask_tokens = objective.patch_pos_embed(mask_tokens)\n",
+    "\n",
+    "            z_hat_proj = objective.feedback_projection(tf_logits)\n",
+    "            ar_context = torch.cat([z_hat_proj[:, :-1], action_summary[:, 1:]], dim=-2)\n",
+    "            ar_logits = objective.heads.get(foresight_cam)(\n",
+    "                CrossAttentionDecoderHead.Input(\n",
+    "                    query=mask_tokens[:, :-1], context=ar_context\n",
+    "                )\n",
+    "            )\n",
+    "            all_ar_preds.append(ar_logits.cpu())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "10185d26023b46108eb7d9f57d49d2b3",
+   "metadata": {},
+   "source": [
+    "## Sanity Check 1 — Rollout Error Curve\n",
+    "\n",
+    "L1 error per prediction step: TF predicts `z[t+1]` from observed `z[t]`; AR predicts `z[t+2]` from its own prediction `z_hat[t+1]`. Both are conditioned on the true action.\n",
+    "\n",
+    "We expect:\n",
+    "- TF step-1 error is the training loss baseline.\n",
+    "- AR step-2 error should be lower (or at least not diverge badly) compared to a model trained without the AR loss."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8763a12b2bbd4a93a75aff182afb95dc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tf_stack = torch.cat(all_tf_preds, dim=0)  # [N_total, T-1, n_patches, D]\n",
+    "tgt_stack = torch.cat(all_targets, dim=0)  # [N_total, T-1, n_patches, D]\n",
+    "\n",
+    "# TF error per timestep slot (mean over batch, patches, dim)\n",
+    "tf_err_per_step = (tf_stack - tgt_stack).abs().mean(dim=(0, 2, 3))  # [T-1]\n",
+    "\n",
+    "fig, ax = plt.subplots(figsize=(7, 4))\n",
+    "steps = np.arange(1, len(tf_err_per_step) + 1)\n",
+    "ax.plot(steps, tf_err_per_step.numpy(), marker=\"o\", label=\"TF (step 1 ... T-1)\")\n",
+    "\n",
+    "if all_ar_preds:\n",
+    "    ar_stack = torch.cat(all_ar_preds, dim=0)  # [N_total, T-2, n_patches, D]\n",
+    "    ar_err_per_step = (ar_stack - tgt_stack[:, 1:]).abs().mean(dim=(0, 2, 3))  # [T-2]\n",
+    "    ar_steps_x = np.arange(2, len(ar_err_per_step) + 2)\n",
+    "    ax.plot(\n",
+    "        ar_steps_x,\n",
+    "        ar_err_per_step.numpy(),\n",
+    "        marker=\"s\",\n",
+    "        linestyle=\"--\",\n",
+    "        label=\"AR (step 2 ... T-1)\",\n",
+    "    )\n",
+    "\n",
+    "ax.set_xlabel(\"Prediction step (1 = next frame)\")\n",
+    "ax.set_ylabel(\"Mean L1 error\")\n",
+    "ax.set_title(\"Rollout error curve — TF vs AR\")\n",
+    "ax.legend()\n",
+    "ax.grid(alpha=0.3)\n",
+    "plt.tight_layout()\n",
+    "plt.savefig(\"rollout_error_curve.png\", dpi=150)\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7623eae2785240b9bd12b16a66d81610",
+   "metadata": {},
+   "source": [
+    "**What to look for:**\n",
+    "- The AR curve at step 2 should be *below* what you'd get if you just naively extended the TF step-1 error (i.e. below the TF value at step 2 from a TF-only model).\n",
+    "- If AR error is *much higher* than TF at the same step, the feedback signal is not helping — revisit `feedback_projection` capacity or learning rate."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7cdc8c89c7104fffa095e18ddfef8986",
+   "metadata": {},
+   "source": [
+    "## Sanity Check 3 — Feature Distribution Overlap\n",
+    "\n",
+    "PCA of predicted patch embeddings vs ground-truth patch embeddings.  \n",
+    "Good AR predictions should overlap with the true distribution rather than drifting.\n",
+    "\n",
+    "We use the first 2 principal components computed on the true embeddings."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b118ea5561624da68c537baed56e602f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def flatten_step(tensor: torch.Tensor, step_idx: int = 0) -> np.ndarray:\n",
+    "    \"\"\"Extract patches at a given timestep: [B, T, N, D] -> [(B*N), D].\"\"\"\n",
+    "    x = tensor[:, step_idx]  # [B, N, D]\n",
+    "    return x.reshape(-1, x.shape[-1]).numpy()\n",
+    "\n",
+    "\n",
+    "true_feats = flatten_step(tgt_stack, step_idx=0)\n",
+    "tf_feats = flatten_step(tf_stack, step_idx=0)\n",
+    "\n",
+    "pca = PCA(n_components=2)\n",
+    "pca.fit(true_feats)\n",
+    "\n",
+    "true_2d = pca.transform(true_feats)\n",
+    "tf_2d = pca.transform(tf_feats)\n",
+    "\n",
+    "n_plots = 2 if all_ar_preds else 1\n",
+    "fig, axes = plt.subplots(1, n_plots, figsize=(6 * n_plots, 5))\n",
+    "if n_plots == 1:\n",
+    "    axes = [axes]\n",
+    "\n",
+    "\n",
+    "def scatter_pair(ax: plt.Axes, pred_2d: np.ndarray, label: str, title: str) -> None:\n",
+    "    ax.scatter(*true_2d.T, s=4, alpha=0.3, label=\"true z\", color=\"steelblue\")\n",
+    "    ax.scatter(*pred_2d.T, s=4, alpha=0.3, label=label, color=\"tomato\")\n",
+    "    ax.set_title(title)\n",
+    "    ax.legend(markerscale=3)\n",
+    "    ax.set_xlabel(\"PC1\")\n",
+    "    ax.set_ylabel(\"PC2\")\n",
+    "\n",
+    "\n",
+    "scatter_pair(axes[0], tf_2d, \"TF pred\", \"TF: predicted vs true (step 1)\")\n",
+    "\n",
+    "if all_ar_preds:\n",
+    "    ar_stack = torch.cat(all_ar_preds, dim=0)\n",
+    "    ar_feats = flatten_step(ar_stack, step_idx=0)\n",
+    "    true_step2_feats = flatten_step(tgt_stack, step_idx=1)\n",
+    "    true_step2_2d = pca.transform(true_step2_feats)\n",
+    "    ar_2d = pca.transform(ar_feats)\n",
+    "\n",
+    "    ax2 = axes[1]\n",
+    "    ax2.scatter(*true_step2_2d.T, s=4, alpha=0.3, label=\"true z[2]\", color=\"steelblue\")\n",
+    "    ax2.scatter(*ar_2d.T, s=4, alpha=0.3, label=\"AR pred z[2]\", color=\"tomato\")\n",
+    "    ax2.set_title(\"AR: predicted vs true (step 2)\")\n",
+    "    ax2.legend(markerscale=3)\n",
+    "    ax2.set_xlabel(\"PC1\")\n",
+    "    ax2.set_ylabel(\"PC2\")\n",
+    "\n",
+    "plt.suptitle(\"Feature distribution overlap — PCA of patch embeddings\", y=1.01)\n",
+    "plt.tight_layout()\n",
+    "plt.savefig(\"feature_distribution.png\", dpi=150)\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "938c804e27f84196a10c8828c723f798",
+   "metadata": {},
+   "source": [
+    "**What to look for:**\n",
+    "- TF and AR predicted distributions should visually overlap with the true distribution (not form a separate cluster).\n",
+    "- If predictions form a tight cluster far from the true distribution, the model is predicting a constant or collapsed mode — check the loss weight and whether the `feedback_projection` has been initialized properly.\n",
+    "- PCA explains variance; for a quantitative score use average cosine similarity between predicted and nearest true embedding."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "504fb2a444614c0babb325280ed9130a",
+   "metadata": {},
+   "source": [
+    "## Bonus: mean cosine similarity"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "59bbdb311c014d738909a11f9e486628",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "true_t = tgt_stack[:, 0].reshape(-1, tgt_stack.shape[-1])\n",
+    "tf_t = tf_stack[:, 0].reshape(-1, tf_stack.shape[-1])\n",
+    "\n",
+    "cos_tf = F.cosine_similarity(tf_t, true_t, dim=-1).mean().item()\n",
+    "print(f\"TF  step-1 mean cosine sim: {cos_tf:.4f}\")  # noqa: T201\n",
+    "\n",
+    "if all_ar_preds:\n",
+    "    true_t2 = tgt_stack[:, 1].reshape(-1, tgt_stack.shape[-1])\n",
+    "    ar_t2 = torch.cat(all_ar_preds, dim=0)[:, 0].reshape(-1, tgt_stack.shape[-1])\n",
+    "    cos_ar = F.cosine_similarity(ar_t2, true_t2, dim=-1).mean().item()\n",
+    "    print(f\"AR  step-2 mean cosine sim: {cos_ar:.4f}\")  # noqa: T201"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

pyproject.toml

@@ -52,6 +52,7 @@ test = ["pytest>=9.0.2", "pytest-lazy-fixtures>=1.4.0"]
 train = [
   "contextily>=1.6.2",
   "matplotlib>=3.10.3",
+  "scikit-learn>=1.6",
   "torchmetrics>=1.9",
   "wandb>=0.24.2",
 ]
@@ -88,6 +89,7 @@ build-backend = "hatchling.build"
 DEP001 = ["app", "rmind"]
 DEP002 = [
   "funcy",
+  "scikit-learn",
   "torchmetrics",
   "pytest",
   "pytest-lazy-fixtures",
@@ -98,6 +100,7 @@ DEP003 = ["rmind"]
 
 [tool.deptry.package_module_name_map]
 hydra-core = "hydra"
+scikit-learn = "sklearn"
 tensordict-nightly = "tensordict"
 pytorch-lightning = [
   "pytorch_lightning",