Initial upload: JEPA-Anything full project + domain checkpoints

01_Vision_Controlled_Binding/README.md

@@ -0,0 +1,12 @@
+# Vision and Controlled Binding
+
+This domain asks whether a visual representation can preserve the structure of a controlled change, not just the appearance of an image. Visible patches provide context and masked patch states provide predictive targets. After JEPA pretraining, a structured readout compares a source scene with its intervention and identifies both the affected spatial support and the operation that occurred. OPF is used to separate complementary visual change modes so the representation can recombine them for support-operation pairs that were not used to fit the readout.
+
+Open a model page for its research background, checkpoint, input format, and inference command. Additional checkpoints from this domain will be released progressively.
+
+## Models
+
+- [Shapes3D_JEPA_Pretraining](Shapes3D_JEPA_Pretraining/README.md)
+- [Shapes3D_Main_Factor_Prediction](Shapes3D_Main_Factor_Prediction/README.md)
+
+[Back to the project overview](../README.md)

01_Vision_Controlled_Binding/Shapes3D_JEPA_Pretraining/README.md

@@ -0,0 +1,36 @@
+# Shapes3D JEPA Pretraining
+
+## Background
+
+This model learns token-level visual representations from controlled changes in object shape, color, scale, orientation, and scene appearance. The archived checkpoint provides the encoder used before factor binding.
+
+## Files
+
+- Model implementation: [`../../model`](<../../model>)
+- Checkpoint: [`checkpoints/Shapes3D_JEPA_Pretraining_final_checkpoint.pt`](checkpoints/Shapes3D_JEPA_Pretraining_final_checkpoint.pt)
+- Inference command: [`../../test/infer_vision.py`](<../../test/infer_vision.py>)
+
+## Run inference
+
+Run from this experiment directory:
+
+```bash
+python ../../test/infer_vision.py \
+  --checkpoint checkpoints/Shapes3D_JEPA_Pretraining_final_checkpoint.pt \
+  --input <preprocessed_images.npz> \
+  --output <new_output_dir>/tokens.npz
+```
+
+Provide an `images` array with shape `[N, 64, 64, 3]` or `[N, 3, 64, 64]`. Values may be uint8 in `[0, 255]` or floats in `[0, 1]`. The command writes visual token embeddings. The Shapes3D release will progressively add the task-specific SO-OPF readout checkpoint alongside the encoder checkpoints.
+
+For a quick checkpoint and tensor-shape check, omit `--input`:
+
+```bash
+python ../../test/infer_vision.py \
+  --checkpoint checkpoints/Shapes3D_JEPA_Pretraining_final_checkpoint.pt \
+  --output /tmp/jepa_inference_smoke/tokens.npz
+```
+
+This command reports the predicted tensor shape and finite-output status. Dataset runs also report MSE when targets are supplied.
+
+[Back to the project overview](<../../README.md>)

01_Vision_Controlled_Binding/Shapes3D_JEPA_Pretraining/checkpoints/Shapes3D_JEPA_Pretraining_final_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:74138fc4a4b537d89ac902e89ca140fc99159e2d1e0c6f883d3ca7952a1dba7d
+size 13441074

01_Vision_Controlled_Binding/Shapes3D_Main_Factor_Prediction/README.md

@@ -0,0 +1,36 @@
+# Shapes3D Factor Prediction
+
+## Background
+
+This stage connects the Shapes3D visual encoder to support-by-operation factor prediction. It asks whether a representation can identify what changed and where the change occurred.
+
+## Files
+
+- Model implementation: [`../../model`](<../../model>)
+- Checkpoint: [`checkpoints/Shapes3D_Main_Factor_Prediction_final_checkpoint.pt`](checkpoints/Shapes3D_Main_Factor_Prediction_final_checkpoint.pt)
+- Inference command: [`../../test/infer_vision.py`](<../../test/infer_vision.py>)
+
+## Run inference
+
+Run from this experiment directory:
+
+```bash
+python ../../test/infer_vision.py \
+  --checkpoint checkpoints/Shapes3D_Main_Factor_Prediction_final_checkpoint.pt \
+  --input <preprocessed_images.npz> \
+  --output <new_output_dir>/tokens.npz
+```
+
+Provide an `images` array with shape `[N, 64, 64, 3]` or `[N, 3, 64, 64]`. Values may be uint8 in `[0, 255]` or floats in `[0, 1]`. The command writes visual token embeddings. The Shapes3D release will progressively add the task-specific SO-OPF readout checkpoint alongside the encoder checkpoints.
+
+For a quick checkpoint and tensor-shape check, omit `--input`:
+
+```bash
+python ../../test/infer_vision.py \
+  --checkpoint checkpoints/Shapes3D_Main_Factor_Prediction_final_checkpoint.pt \
+  --output /tmp/jepa_inference_smoke/tokens.npz
+```
+
+This command reports the predicted tensor shape and finite-output status. Dataset runs also report MSE when targets are supplied.
+
+[Back to the project overview](<../../README.md>)

01_Vision_Controlled_Binding/Shapes3D_Main_Factor_Prediction/checkpoints/Shapes3D_Main_Factor_Prediction_final_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb45ceb88443d21d3f62e95b8a5e9495cb30b5d04e8aa9725954caecbca020e7
+size 11157882

02_Single_Cell/Norman_Single_Cell_Perturbation_Prediction/README.md

@@ -0,0 +1,44 @@
+# Norman Single-Cell Perturbation Prediction
+
+## Background
+
+This model predicts how gene expression changes after a perturbation. It combines a single-cell backbone with factorized perturbation-response representations, then decodes the predicted expression profile.
+
+## Files and requirements
+
+- Model code: [`model/cell_jepa`](model/cell_jepa)
+- Vocabulary: [`model/Norman_Single_Cell_Perturbation_Prediction_vocab.json`](model/Norman_Single_Cell_Perturbation_Prediction_vocab.json)
+- Checkpoint: [`checkpoints/Norman_Single_Cell_Perturbation_Prediction_final_checkpoint.pt`](checkpoints/Norman_Single_Cell_Perturbation_Prediction_final_checkpoint.pt)
+- Inference and metric entry point: [`test/Norman_Single_Cell_Perturbation_Prediction_test.py`](test/Norman_Single_Cell_Perturbation_Prediction_test.py)
+
+The model uses scGPT for the expression backbone and GEARS for the Norman evaluation pipeline.
+
+## Verify checkpoint loading
+
+Use checkpoint mode to reconstruct the model and report its parameter and factor configuration:
+
+```bash
+python test/Norman_Single_Cell_Perturbation_Prediction_test.py \
+  --repo <cell_jepa_source_repo> \
+  --checkpoint checkpoints/Norman_Single_Cell_Perturbation_Prediction_final_checkpoint.pt \
+  --vocab model/Norman_Single_Cell_Perturbation_Prediction_vocab.json \
+  --checkpoint-only \
+  --output <new_output_dir>/norman_checkpoint.json \
+  --device cpu
+```
+
+## Evaluate on Norman
+
+```bash
+python test/Norman_Single_Cell_Perturbation_Prediction_test.py \
+  --repo <cell_jepa_source_repo> \
+  --data-root <norman_data_root> \
+  --checkpoint checkpoints/Norman_Single_Cell_Perturbation_Prediction_final_checkpoint.pt \
+  --vocab model/Norman_Single_Cell_Perturbation_Prediction_vocab.json \
+  --output <new_output_dir>/norman_metrics.json \
+  --device cuda
+```
+
+The evaluation command predicts expression profiles on the prepared Norman split and writes the perturbation metrics to JSON.
+
+[Back to the project overview](../../README.md)

02_Single_Cell/Norman_Single_Cell_Perturbation_Prediction/checkpoints/Norman_Single_Cell_Perturbation_Prediction_final_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ee029201f537faea600ec9d8f7c5dfad9b0a7eaac747eb9bc60b407f70f3dcb7
+size 419246117

02_Single_Cell/Norman_Single_Cell_Perturbation_Prediction/model/cell_jepa/__init__.py

@@ -0,0 +1,12 @@
+"""Cell-JEPA modules used by the reproduction workflow."""
+
+from .objectives import CellJEPA, PerturbationJEPA, supervised_ecs_loss
+from .objectives_anything import CellJEPAAnything, PerturbationJEPAAnything
+
+__all__ = [
+    "CellJEPA",
+    "CellJEPAAnything",
+    "PerturbationJEPA",
+    "PerturbationJEPAAnything",
+    "supervised_ecs_loss",
+]

02_Single_Cell/Norman_Single_Cell_Perturbation_Prediction/model/cell_jepa/objectives.py

@@ -0,0 +1,220 @@
+"""JEPA losses wrapped around the official scGPT model implementations.
+
+The paper's public artifact is a manuscript, rather than an official code
+repository.  This module implements Eqs. 2.3 and 2.5 using the public scGPT
+backbone and GEARS perturbation data protocol.
+"""
+
+from __future__ import annotations
+
+import copy
+from typing import Any, Optional
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+
+def _mlp(width: int) -> nn.Sequential:
+    return nn.Sequential(nn.Linear(width, width), nn.GELU(), nn.Linear(width, width))
+
+
+def supervised_ecs_loss(embeddings: Tensor, labels: Tensor, temperature: float = 0.5) -> Tensor:
+    """Label-aware ECS used for PBMC representation learning (paper Sec. 2.4.2).
+
+    Each cell treats all same-label cells in its minibatch as positives. Anchors
+    without another same-label cell are ignored rather than producing NaNs.
+    """
+    if embeddings.ndim != 2 or labels.ndim != 1:
+        raise ValueError("embeddings must be [batch, width] and labels must be [batch]")
+    if embeddings.shape[0] != labels.shape[0]:
+        raise ValueError("labels and embeddings must have the same batch dimension")
+    if embeddings.shape[0] < 2:
+        return embeddings.new_zeros(())
+    logits = F.normalize(embeddings, dim=-1) @ F.normalize(embeddings, dim=-1).T
+    logits = logits / temperature
+    diagonal = torch.eye(len(labels), dtype=torch.bool, device=labels.device)
+    logits = logits.masked_fill(diagonal, float("-inf"))
+    log_prob = F.log_softmax(logits, dim=1)
+    positive = labels[:, None].eq(labels[None, :]) & ~diagonal
+    count = positive.sum(dim=1)
+    valid = count > 0
+    if not valid.any():
+        return embeddings.new_zeros(())
+    return -((log_prob.masked_fill(~positive, 0).sum(dim=1) / count.clamp_min(1))[valid]).mean()
+
+
+def ecs_loss(embeddings: Tensor, threshold: float = 0.3) -> Tensor:
+    """The original scGPT embedding-consistency penalty used for perturbation.
+
+    Perturbation batches contain a single cell line rather than PBMC cell-type
+    labels, so Sec. 2.5 uses the non-label-aware ECS term.  This is equivalent
+    to the public scGPT implementation but is exposed here because the JEPA
+    wrapper needs the intermediate cell embeddings.
+    """
+    if embeddings.shape[0] < 2:
+        return embeddings.new_zeros(())
+    similarity = F.normalize(embeddings, dim=-1) @ F.normalize(embeddings, dim=-1).T
+    similarity = similarity - torch.diag_embed(torch.diagonal(similarity))
+    similarity = F.relu(similarity)
+    return torch.mean(1 - (similarity - threshold).pow(2))
+
+
+class CellJEPA(nn.Module):
+    """Student--EMA teacher wrapper for ``scgpt.model.TransformerModel``.
+
+    Student inputs may contain a mask sentinel; the teacher always receives
+    the corresponding unmasked values. ``losses`` returns masked-value MSE
+    plus the latent cosine-distance JEPA objective from Eq. 2.3.
+    """
+
+    def __init__(self, student: nn.Module, momentum: float = 0.996) -> None:
+        super().__init__()
+        if not 0.0 <= momentum < 1.0:
+            raise ValueError("momentum must be in [0, 1)")
+        self.student = student
+        self.teacher = copy.deepcopy(student)
+        self.predictor = _mlp(student.d_model)
+        self.momentum = momentum
+        self.freeze_teacher()
+
+    def freeze_teacher(self) -> None:
+        self.teacher.eval()
+        for parameter in self.teacher.parameters():
+            parameter.requires_grad_(False)
+
+    def train(self, mode: bool = True) -> "CellJEPA":
+        super().train(mode)
+        self.teacher.eval()  # the EMA target must never use dropout updates
+        return self
+
+    @torch.no_grad()
+    def update_teacher(self) -> None:
+        """EMA update performed once after each student optimizer step."""
+        for teacher_parameter, student_parameter in zip(
+            self.teacher.parameters(), self.student.parameters()
+        ):
+            teacher_parameter.mul_(self.momentum).add_(student_parameter, alpha=1 - self.momentum)
+        for teacher_buffer, student_buffer in zip(self.teacher.buffers(), self.student.buffers()):
+            teacher_buffer.copy_(student_buffer)
+
+    def forward(
+        self,
+        gene_ids: Tensor,
+        masked_values: Tensor,
+        target_values: Tensor,
+        padding_mask: Tensor,
+        batch_labels: Optional[Tensor] = None,
+        include_gepc: bool = False,
+    ) -> tuple[dict[str, Tensor], Tensor]:
+        student_output = self.student(
+            gene_ids, masked_values, padding_mask, batch_labels=batch_labels,
+            CLS=False, CCE=False, MVC=include_gepc, ECS=False,
+        )
+        with torch.no_grad():
+            teacher_output = self.teacher(
+                gene_ids, target_values, padding_mask, batch_labels=batch_labels,
+                CLS=False, CCE=False, MVC=False, ECS=False,
+            )
+        return student_output, teacher_output["cell_emb"]
+
+    def losses(
+        self,
+        gene_ids: Tensor,
+        masked_values: Tensor,
+        target_values: Tensor,
+        padding_mask: Tensor,
+        masked_positions: Tensor,
+        batch_labels: Optional[Tensor] = None,
+        cell_labels: Optional[Tensor] = None,
+        jepa_weight: float = 1.0,
+        reconstruction_weight: float = 1.0,
+        ecs_weight: float = 0.0,
+        gepc_weight: float = 0.0,
+    ) -> dict[str, Tensor]:
+        student_output, teacher_embedding = self(
+            gene_ids, masked_values, target_values, padding_mask, batch_labels,
+            include_gepc=bool(gepc_weight),
+        )
+        prediction = student_output["mlm_output"]
+        reconstruction = F.mse_loss(prediction[masked_positions], target_values[masked_positions])
+        jepa = 1 - F.cosine_similarity(self.predictor(student_output["cell_emb"]), teacher_embedding, dim=-1).mean()
+        ecs = (
+            supervised_ecs_loss(student_output["cell_emb"], cell_labels)
+            if cell_labels is not None and ecs_weight
+            else reconstruction.new_zeros(())
+        )
+        gepc = (
+            F.mse_loss(student_output["mvc_output"][masked_positions], target_values[masked_positions])
+            if gepc_weight
+            else reconstruction.new_zeros(())
+        )
+        total = reconstruction_weight * reconstruction + jepa_weight * jepa + ecs_weight * ecs + gepc_weight * gepc
+        return {"total": total, "reconstruction": reconstruction, "gepc": gepc, "jepa": jepa, "ecs": ecs}
+
+
+class PerturbationJEPA(nn.Module):
+    """Cell-JEPA perturbation head for scGPT's ``TransformerGenerator``.
+
+    The student receives control expression and perturbation indicators. The
+    EMA teacher receives the observed post-perturbation profile and the same
+    perturbation indicators. This follows paper Sec. 2.5 with no input masking.
+    """
+
+    def __init__(self, student: nn.Module, momentum: float = 0.996) -> None:
+        super().__init__()
+        self.student = student
+        self.teacher = copy.deepcopy(student)
+        self.predictor = _mlp(student.d_model)
+        self.momentum = momentum
+        self.freeze_teacher()
+
+    def freeze_teacher(self) -> None:
+        self.teacher.eval()
+        for parameter in self.teacher.parameters():
+            parameter.requires_grad_(False)
+
+    def train(self, mode: bool = True) -> "PerturbationJEPA":
+        super().train(mode)
+        self.teacher.eval()
+        return self
+
+    @torch.no_grad()
+    def update_teacher(self) -> None:
+        for teacher_parameter, student_parameter in zip(self.teacher.parameters(), self.student.parameters()):
+            teacher_parameter.mul_(self.momentum).add_(student_parameter, alpha=1 - self.momentum)
+        for teacher_buffer, student_buffer in zip(self.teacher.buffers(), self.student.buffers()):
+            teacher_buffer.copy_(student_buffer)
+
+    @staticmethod
+    def _run(model: nn.Module, gene_ids: Tensor, values: Tensor, perturbations: Tensor, padding_mask: Tensor) -> tuple[Tensor, Tensor]:
+        # The public TransformerGenerator keeps these intermediates private;
+        # accessing them avoids modifying the official baseline source.
+        hidden = model._encode(gene_ids, values, perturbations, padding_mask)
+        decoded = model.decoder(hidden, values)["pred"]
+        embedding = model._get_cell_emb_from_layer(hidden, values)
+        return decoded, embedding
+
+    def losses(
+        self,
+        gene_ids: Tensor,
+        control_values: Tensor,
+        perturbations: Tensor,
+        observed_values: Tensor,
+        padding_mask: Tensor,
+        jepa_weight: float = 1.0,
+        reconstruction_weight: float = 1.0,
+        ecs_weight: float = 0.8,
+    ) -> dict[str, Tensor]:
+        predicted_values, student_embedding = self._run(
+            self.student, gene_ids, control_values, perturbations, padding_mask
+        )
+        with torch.no_grad():
+            _, teacher_embedding = self._run(
+                self.teacher, gene_ids, observed_values, perturbations, padding_mask
+            )
+        reconstruction = F.mse_loss(predicted_values, observed_values)
+        jepa = 1 - F.cosine_similarity(self.predictor(student_embedding), teacher_embedding, dim=-1).mean()
+        consistency = ecs_loss(student_embedding) if ecs_weight else reconstruction.new_zeros(())
+        total = reconstruction_weight * reconstruction + jepa_weight * jepa + ecs_weight * consistency
+        return {"total": total, "reconstruction": reconstruction, "jepa": jepa, "ecs": consistency}

02_Single_Cell/Norman_Single_Cell_Perturbation_Prediction/model/cell_jepa/objectives_anything.py

@@ -0,0 +1,558 @@
+"""JEPA-Anything with response factors and activation-level diversity."""
+
+from __future__ import annotations
+
+import copy
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+from .objectives import ecs_loss, supervised_ecs_loss
+
+
+class PerturbationJEPAAnything(nn.Module):
+    """Predict orthogonal perturbation-response factors and use them for decoding.
+
+    The objective targets the paired response
+    ``teacher(observed) - teacher(control)``, penalizes cross-factor
+    activation covariance, learns sparse balanced routing, and injects the
+    routed factor prediction into the expression decoder.
+    """
+
+    def __init__(
+        self,
+        student: nn.Module,
+        num_factors: int = 4,
+        momentum: float = 0.996,
+    ) -> None:
+        super().__init__()
+        width = int(student.d_model)
+        if num_factors < 2 or width % num_factors:
+            raise ValueError("num_factors must be >=2 and divide student.d_model")
+        if not 0.0 <= momentum < 1.0:
+            raise ValueError("momentum must be in [0, 1)")
+
+        self.student = student
+        self.teacher = copy.deepcopy(student)
+        self.width = width
+        self.num_factors = num_factors
+        self.factor_width = width // num_factors
+        self.momentum = momentum
+
+        identity = torch.eye(width).reshape(width, num_factors, self.factor_width)
+        self.projectors = nn.Parameter(identity.permute(1, 0, 2).contiguous())
+        self.predictor_trunk = nn.Sequential(nn.Linear(width, width), nn.GELU())
+        self.factor_heads = nn.ModuleList(
+            nn.Linear(width, self.factor_width) for _ in range(num_factors)
+        )
+        self.gate = nn.Linear(width, num_factors)
+        self.response_to_hidden = nn.Linear(width, width)
+        nn.init.zeros_(self.response_to_hidden.weight)
+        nn.init.zeros_(self.response_to_hidden.bias)
+        self.freeze_teacher()
+
+    def freeze_teacher(self) -> None:
+        self.teacher.eval()
+        for parameter in self.teacher.parameters():
+            parameter.requires_grad_(False)
+
+    def train(self, mode: bool = True) -> "PerturbationJEPAAnything":
+        super().train(mode)
+        self.teacher.eval()
+        return self
+
+    @torch.no_grad()
+    def update_teacher(self) -> None:
+        for teacher_parameter, student_parameter in zip(
+            self.teacher.parameters(), self.student.parameters()
+        ):
+            teacher_parameter.mul_(self.momentum).add_(
+                student_parameter, alpha=1 - self.momentum
+            )
+        for teacher_buffer, student_buffer in zip(
+            self.teacher.buffers(), self.student.buffers()
+        ):
+            teacher_buffer.copy_(student_buffer)
+
+    @torch.no_grad()
+    def retract_projectors(self) -> None:
+        """Project the concatenated blocks back onto the orthogonal group."""
+        full = torch.cat(tuple(self.projectors), dim=1)
+        q, r = torch.linalg.qr(full)
+        signs = torch.sign(torch.diagonal(r))
+        signs[signs == 0] = 1
+        q = q * signs.unsqueeze(0)
+        blocks = q.split(self.factor_width, dim=1)
+        self.projectors.copy_(torch.stack(blocks, dim=0))
+
+    @staticmethod
+    def _encode(
+        model: nn.Module,
+        gene_ids: Tensor,
+        values: Tensor,
+        perturbations: Tensor,
+        padding_mask: Tensor,
+    ) -> tuple[Tensor, Tensor]:
+        hidden = model._encode(gene_ids, values, perturbations, padding_mask)
+        embedding = model._get_cell_emb_from_layer(hidden, values)
+        return hidden, embedding
+
+    def _predict_from_hidden(
+        self,
+        hidden: Tensor,
+        embedding: Tensor,
+        control_values: Tensor,
+    ) -> tuple[Tensor, Tensor, Tensor]:
+        shared = self.predictor_trunk(embedding)
+        predicted_factors = torch.stack(
+            [head(shared) for head in self.factor_heads], dim=1
+        )
+        gates = torch.softmax(self.gate(embedding), dim=1)
+        routed = predicted_factors * gates.unsqueeze(-1) * self.num_factors
+        response_context = self.response_to_hidden(routed.flatten(1))
+        predicted_values = self.student.decoder(
+            hidden + response_context.unsqueeze(1), control_values
+        )["pred"]
+        return predicted_values, predicted_factors, gates
+
+    def predict_expression(
+        self,
+        gene_ids: Tensor,
+        control_values: Tensor,
+        perturbations: Tensor,
+        padding_mask: Tensor,
+    ) -> Tensor:
+        hidden, embedding = self._encode(
+            self.student, gene_ids, control_values, perturbations, padding_mask
+        )
+        predicted_values, _, _ = self._predict_from_hidden(
+            hidden, embedding, control_values
+        )
+        return predicted_values
+
+    def _block_orthogonality_loss(self) -> Tensor:
+        identity = torch.eye(
+            self.factor_width,
+            device=self.projectors.device,
+            dtype=self.projectors.dtype,
+        )
+        within = torch.stack(
+            [
+                (projector.T @ projector - identity).pow(2).mean()
+                for projector in self.projectors
+            ]
+        ).mean()
+        cross_terms = []
+        for i in range(self.num_factors):
+            for j in range(i + 1, self.num_factors):
+                cross_terms.append(
+                    (self.projectors[i].T @ self.projectors[j]).pow(2).mean()
+                )
+        cross = torch.stack(cross_terms).mean()
+        return within + cross
+
+    @staticmethod
+    def _cross_factor_covariance_loss(factors: Tensor) -> Tensor:
+        """VICReg-style cross-covariance between different factor blocks."""
+        batch_size = factors.shape[0]
+        centered = factors - factors.mean(dim=0, keepdim=True)
+        std = centered.std(dim=0, unbiased=False, keepdim=True).clamp_min(1e-4)
+        normalized = centered / std
+        covariance = torch.einsum("bkr,bls->klrs", normalized, normalized)
+        covariance = covariance / max(batch_size - 1, 1)
+        n_factors = factors.shape[1]
+        mask = ~torch.eye(n_factors, dtype=torch.bool, device=factors.device)
+        return covariance[mask].pow(2).mean()
+
+    def losses(
+        self,
+        gene_ids: Tensor,
+        control_values: Tensor,
+        perturbations: Tensor,
+        observed_values: Tensor,
+        padding_mask: Tensor,
+        reconstruction_weight: float = 1.0,
+        prediction_weight: float = 1.0,
+        ecs_weight: float = 0.8,
+        orthogonality_weight: float = 1.0,
+        cross_covariance_weight: float = 1.0,
+        factor_activity_weight: float = 1.0,
+        encoder_variance_weight: float = 1.0,
+        gate_entropy_weight: float = 0.05,
+        gate_balance_weight: float = 0.05,
+        factor_std_floor: float = 0.02,
+        encoder_std_floor: float = 0.1,
+    ) -> dict[str, Tensor]:
+        student_hidden, student_embedding = self._encode(
+            self.student, gene_ids, control_values, perturbations, padding_mask
+        )
+        predicted_values, predicted_factors, gates = self._predict_from_hidden(
+            student_hidden, student_embedding, control_values
+        )
+        # The JEPA target is a difference between two nearby teacher
+        # embeddings.  Computing either embedding in BF16 causes cancellation
+        # and quantization noise in that small response, so the teacher branch
+        # deliberately remains FP32 even when the student uses autocast.
+        with torch.no_grad(), torch.autocast(
+            device_type=control_values.device.type, enabled=False
+        ):
+            _, teacher_control = self._encode(
+                self.teacher,
+                gene_ids,
+                control_values.float(),
+                perturbations,
+                padding_mask,
+            )
+            _, teacher_observed = self._encode(
+                self.teacher,
+                gene_ids,
+                observed_values.float(),
+                perturbations,
+                padding_mask,
+            )
+            teacher_response = teacher_observed - teacher_control
+
+        # Keep the expensive transformer/decoder path under the caller's
+        # autocast, but compute all covariance, variance, routing, and scalar
+        # losses in FP32.  These statistics are sensitive to BF16 rounding and
+        # otherwise drift even when the forward activations remain stable.
+        with torch.autocast(device_type=control_values.device.type, enabled=False):
+            predicted_values_fp32 = predicted_values.float()
+            predicted_factors_fp32 = predicted_factors.float()
+            observed_values_fp32 = observed_values.float()
+            student_embedding_fp32 = student_embedding.float()
+            gates_fp32 = gates.float()
+            teacher_response_fp32 = teacher_response.detach().float()
+            target_factors = torch.einsum(
+                "bd,kdr->bkr", teacher_response_fp32, self.projectors.float()
+            )
+            reconstruction = F.mse_loss(
+                predicted_values_fp32, observed_values_fp32
+            )
+            prediction = F.mse_loss(predicted_factors_fp32, target_factors)
+            orthogonality = self._block_orthogonality_loss()
+            cross_covariance = self._cross_factor_covariance_loss(target_factors)
+            factor_std = target_factors.std(dim=0, unbiased=False)
+            encoder_std = student_embedding_fp32.std(dim=0, unbiased=False)
+            factor_activity = F.relu(factor_std_floor - factor_std).mean()
+            encoder_variance = F.relu(encoder_std_floor - encoder_std).mean()
+            consistency = (
+                ecs_loss(student_embedding_fp32)
+                if ecs_weight
+                else reconstruction.new_zeros(())
+            )
+            gate_entropy = -(
+                gates_fp32.clamp_min(1e-8)
+                * gates_fp32.clamp_min(1e-8).log()
+            ).sum(dim=1).mean() / math.log(self.num_factors)
+            mean_gate = gates_fp32.mean(dim=0)
+            gate_balance = self.num_factors * (
+                mean_gate - 1 / self.num_factors
+            ).pow(2).sum()
+            effective_factors = torch.exp(
+                -(
+                    gates_fp32.clamp_min(1e-8)
+                    * gates_fp32.clamp_min(1e-8).log()
+                ).sum(dim=1)
+            ).mean()
+
+        total = (
+            reconstruction_weight * reconstruction
+            + prediction_weight * prediction
+            + ecs_weight * consistency
+            + orthogonality_weight * orthogonality
+            + cross_covariance_weight * cross_covariance
+            + factor_activity_weight * factor_activity
+            + encoder_variance_weight * encoder_variance
+            + gate_entropy_weight * gate_entropy
+            + gate_balance_weight * gate_balance
+        )
+        return {
+            "total": total,
+            "reconstruction": reconstruction,
+            "prediction": prediction,
+            "ecs": consistency,
+            "orthogonality": orthogonality,
+            "cross_covariance": cross_covariance,
+            "factor_activity": factor_activity,
+            "encoder_variance": encoder_variance,
+            "gate_entropy": gate_entropy,
+            "gate_balance": gate_balance,
+            "effective_factors": effective_factors,
+            "factor_std_min": factor_std.min(),
+            "factor_std_mean": factor_std.mean(),
+            "encoder_std_min": encoder_std.min(),
+            "encoder_std_mean": encoder_std.mean(),
+        }
+
+    @torch.no_grad()
+    def factor_diagnostics(self) -> dict[str, float]:
+        full = torch.cat(tuple(self.projectors), dim=1).float()
+        gram = full.T @ full
+        identity = torch.eye(self.width, device=full.device)
+        singular = torch.linalg.svdvals(full)
+        return {
+            "orthogonality_rmse": float((gram - identity).pow(2).mean().sqrt()),
+            "orthogonality_max_abs": float((gram - identity).abs().max()),
+            "projection_singular_min": float(singular.min()),
+            "projection_singular_max": float(singular.max()),
+            "projection_condition_number": float(singular.max() / singular.min()),
+        }
+
+
+class CellJEPAAnything(nn.Module):
+    """JEPA-Anything objective for cell-type representation learning.
+
+    Unlike the perturbation objective above, PBMC has no observed/control pair.
+    The EMA teacher therefore supplies the target cell representation while the
+    student sees the masked expression profile.  The representation is routed
+    through orthogonal response factors so this mode uses the same
+    JEPA-Anything decomposition and diversity constraints as the perturbation
+    experiments.
+    """
+
+    def __init__(self, student: nn.Module, num_factors: int = 4, momentum: float = 0.996) -> None:
+        super().__init__()
+        width = int(student.d_model)
+        if num_factors < 2 or width % num_factors:
+            raise ValueError("num_factors must be >=2 and divide student.d_model")
+        self.student = student
+        self.teacher = copy.deepcopy(student)
+        self.width = width
+        self.num_factors = num_factors
+        self.factor_width = width // num_factors
+        self.momentum = momentum
+        identity = torch.eye(width).reshape(width, num_factors, self.factor_width)
+        self.projectors = nn.Parameter(identity.permute(1, 0, 2).contiguous())
+        self.predictor_trunk = nn.Sequential(nn.Linear(width, width), nn.GELU())
+        self.factor_heads = nn.ModuleList(nn.Linear(width, self.factor_width) for _ in range(num_factors))
+        self.gate = nn.Linear(width, num_factors)
+        self.freeze_teacher()
+
+    def freeze_teacher(self) -> None:
+        self.teacher.eval()
+        for parameter in self.teacher.parameters():
+            parameter.requires_grad_(False)
+
+    def train(self, mode: bool = True) -> "CellJEPAAnything":
+        super().train(mode)
+        self.teacher.eval()
+        return self
+
+    @torch.no_grad()
+    def update_teacher(self) -> None:
+        for teacher_parameter, student_parameter in zip(self.teacher.parameters(), self.student.parameters()):
+            teacher_parameter.mul_(self.momentum).add_(student_parameter, alpha=1 - self.momentum)
+        for teacher_buffer, student_buffer in zip(self.teacher.buffers(), self.student.buffers()):
+            teacher_buffer.copy_(student_buffer)
+
+    @torch.no_grad()
+    def retract_projectors(self) -> None:
+        full = torch.cat(tuple(self.projectors), dim=1)
+        q, r = torch.linalg.qr(full)
+        signs = torch.sign(torch.diagonal(r)).masked_fill(torch.diagonal(r) == 0, 1)
+        q = q * signs.unsqueeze(0)
+        self.projectors.copy_(torch.stack(q.split(self.factor_width, dim=1), dim=0))
+
+    @staticmethod
+    def _cross_factor_covariance_loss(factors: Tensor) -> Tensor:
+        centered = factors - factors.mean(dim=0, keepdim=True)
+        std = centered.std(dim=0, unbiased=False, keepdim=True).clamp_min(1e-4)
+        normalized = centered / std
+        covariance = torch.einsum("bkr,bls->klrs", normalized, normalized)
+        covariance = covariance / max(factors.shape[0] - 1, 1)
+        mask = ~torch.eye(factors.shape[1], dtype=torch.bool, device=factors.device)
+        return covariance[mask].pow(2).mean()
+
+    def losses(
+        self, gene_ids: Tensor, masked_values: Tensor, target_values: Tensor,
+        padding_mask: Tensor, masked_positions: Tensor, batch_labels: Tensor | None = None,
+        cell_labels: Tensor | None = None, reconstruction_weight: float = 1.0,
+        jepa_weight: float = 1.0, ecs_weight: float = 1.0,
+    ) -> dict[str, Tensor]:
+        # The official scGPT checkpoint used for PBMC does not enable either
+        # batch-label conditioning mode.  Passing labels in that configuration
+        # makes scGPT reject the batch, so only forward them when supported.
+        batch_kwargs = {}
+        if batch_labels is not None and (
+            getattr(self.student, "use_batch_labels", False)
+            or getattr(self.student, "domain_spec_batchnorm", False)
+        ):
+            batch_kwargs["batch_labels"] = batch_labels
+        student_output = self.student(
+            gene_ids, masked_values, padding_mask, **batch_kwargs,
+            CLS=False, CCE=False, MVC=True, ECS=False,
+        )
+        with torch.no_grad():
+            teacher_output = self.teacher(
+                gene_ids, target_values, padding_mask, **batch_kwargs,
+                CLS=False, CCE=False, MVC=False, ECS=False,
+            )
+        student_embedding = student_output["cell_emb"]
+        teacher_embedding = teacher_output["cell_emb"].detach()
+        predicted_factors = torch.stack(
+            [head(self.predictor_trunk(student_embedding)) for head in self.factor_heads], dim=1
+        )
+        gates = torch.softmax(self.gate(student_embedding), dim=1)
+        routed = predicted_factors * gates.unsqueeze(-1) * self.num_factors
+        target_factors = torch.einsum("bd,kdr->bkr", teacher_embedding.float(), self.projectors.float())
+        with torch.autocast(device_type=masked_values.device.type, enabled=False):
+            reconstruction = F.mse_loss(student_output["mlm_output"].float()[masked_positions], target_values.float()[masked_positions])
+            gepc = F.mse_loss(student_output["mvc_output"].float()[masked_positions], target_values.float()[masked_positions])
+            prediction = F.mse_loss(routed.float(), target_factors)
+            identity = torch.eye(self.factor_width, device=self.projectors.device)
+            within = torch.stack([(p.T @ p - identity).pow(2).mean() for p in self.projectors]).mean()
+            cross = []
+            for i in range(self.num_factors):
+                for j in range(i + 1, self.num_factors):
+                    cross.append((self.projectors[i].T @ self.projectors[j]).pow(2).mean())
+            orthogonality = within + torch.stack(cross).mean()
+            cross_covariance = self._cross_factor_covariance_loss(target_factors)
+            factor_std = target_factors.std(dim=0, unbiased=False)
+            encoder_std = student_embedding.float().std(dim=0, unbiased=False)
+            factor_activity = F.relu(0.02 - factor_std).mean()
+            encoder_variance = F.relu(0.1 - encoder_std).mean()
+            ecs = supervised_ecs_loss(student_embedding.float(), cell_labels) if cell_labels is not None else student_embedding.new_zeros(())
+            entropy = -(gates.float().clamp_min(1e-8) * gates.float().clamp_min(1e-8).log()).sum(dim=1).mean() / math.log(self.num_factors)
+            mean_gate = gates.float().mean(dim=0)
+            balance = self.num_factors * (mean_gate - 1 / self.num_factors).pow(2).sum()
+            effective = torch.exp(-(gates.float().clamp_min(1e-8) * gates.float().clamp_min(1e-8).log()).sum(dim=1)).mean()
+        total = reconstruction_weight * reconstruction + jepa_weight * (prediction + 0.5 * gepc) + ecs_weight * ecs + orthogonality + cross_covariance + factor_activity + encoder_variance - 0.05 * entropy + 0.05 * balance
+        return {"total": total, "reconstruction": reconstruction, "gepc": gepc, "prediction": prediction, "ecs": ecs, "orthogonality": orthogonality, "cross_covariance": cross_covariance, "factor_activity": factor_activity, "encoder_variance": encoder_variance, "effective_factors": effective}
+
+    @torch.no_grad()
+    def factor_diagnostics(self) -> dict[str, float]:
+        full = torch.cat(tuple(self.projectors), dim=1).float()
+        gram = full.T @ full
+        singular = torch.linalg.svdvals(full)
+        return {"orthogonality_rmse": float((gram - torch.eye(self.width, device=full.device)).pow(2).mean().sqrt()), "projection_singular_min": float(singular.min()), "projection_singular_max": float(singular.max()), "effective_factors": float(self.num_factors)}
+
+
+class CellJEPAAnythingStrict(nn.Module):
+    """Paper-aligned OPF objective for the PBMC representation experiment.
+
+    This is the direct Eq. (3)--(10) implementation: an unmasked EMA target,
+    one independent predictor per orthogonal factor, and the orthogonality,
+    target-activity, and online-encoder-variance terms.  It intentionally has
+    no routing gate, label-supervised ECS, or observation-space reconstruction;
+    those are separate extensions and are not part of the paper objective.
+    """
+
+    def __init__(self, student: nn.Module, num_factors: int = 8, momentum: float = 0.996) -> None:
+        super().__init__()
+        width = int(student.d_model)
+        if num_factors < 2 or width % num_factors:
+            raise ValueError("num_factors must be >=2 and divide student.d_model")
+        self.student = student
+        self.teacher = copy.deepcopy(student)
+        self.width = width
+        self.num_factors = num_factors
+        self.factor_width = width // num_factors
+        self.momentum = momentum
+        identity = torch.eye(width).reshape(width, num_factors, self.factor_width)
+        self.projectors = nn.Parameter(identity.permute(1, 0, 2).contiguous())
+        self.predictor_trunk = nn.Sequential(nn.Linear(width, width), nn.GELU())
+        self.factor_heads = nn.ModuleList(
+            nn.Linear(width, self.factor_width) for _ in range(num_factors)
+        )
+        self.freeze_teacher()
+
+    def freeze_teacher(self) -> None:
+        self.teacher.eval()
+        for parameter in self.teacher.parameters():
+            parameter.requires_grad_(False)
+
+    def train(self, mode: bool = True) -> "CellJEPAAnythingStrict":
+        super().train(mode)
+        self.teacher.eval()
+        return self
+
+    @torch.no_grad()
+    def update_teacher(self) -> None:
+        for teacher_parameter, student_parameter in zip(
+            self.teacher.parameters(), self.student.parameters()
+        ):
+            teacher_parameter.mul_(self.momentum).add_(
+                student_parameter, alpha=1 - self.momentum
+            )
+        for teacher_buffer, student_buffer in zip(self.teacher.buffers(), self.student.buffers()):
+            teacher_buffer.copy_(student_buffer)
+
+    @torch.no_grad()
+    def retract_projectors(self) -> None:
+        full = torch.cat(tuple(self.projectors), dim=1)
+        q, r = torch.linalg.qr(full)
+        signs = torch.sign(torch.diagonal(r)).masked_fill(torch.diagonal(r) == 0, 1)
+        self.projectors.copy_(torch.stack(q.mul(signs.unsqueeze(0)).split(self.factor_width, dim=1), dim=0))
+
+    def losses(
+        self,
+        gene_ids: Tensor,
+        masked_values: Tensor,
+        target_values: Tensor,
+        padding_mask: Tensor,
+        masked_positions: Tensor | None = None,
+        batch_labels: Tensor | None = None,
+        cell_labels: Tensor | None = None,
+    ) -> dict[str, Tensor]:
+        batch_kwargs = {}
+        if batch_labels is not None and (
+            getattr(self.student, "use_batch_labels", False)
+            or getattr(self.student, "domain_spec_batchnorm", False)
+        ):
+            batch_kwargs["batch_labels"] = batch_labels
+        student_output = self.student(
+            gene_ids, masked_values, padding_mask, **batch_kwargs,
+            CLS=False, CCE=False, MVC=False, ECS=False,
+        )
+        with torch.no_grad():
+            teacher_output = self.teacher(
+                gene_ids, target_values, padding_mask, **batch_kwargs,
+                CLS=False, CCE=False, MVC=False, ECS=False,
+            )
+        student_embedding = student_output["cell_emb"]
+        teacher_embedding = teacher_output["cell_emb"].detach().float()
+        predicted_factors = torch.stack(
+            [head(self.predictor_trunk(student_embedding)) for head in self.factor_heads],
+            dim=1,
+        )
+        target_factors = torch.einsum(
+            "bd,kdr->bkr", teacher_embedding, self.projectors.float()
+        )
+        with torch.autocast(device_type=masked_values.device.type, enabled=False):
+            prediction = F.mse_loss(predicted_factors.float(), target_factors)
+            identity = torch.eye(self.factor_width, device=self.projectors.device)
+            within = torch.stack(
+                [(p.T @ p - identity).pow(2).mean() for p in self.projectors]
+            ).mean()
+            cross_terms = []
+            for i in range(self.num_factors):
+                for j in range(i + 1, self.num_factors):
+                    cross_terms.append((self.projectors[i].T @ self.projectors[j]).pow(2).mean())
+            orthogonality = within + torch.stack(cross_terms).mean()
+            factor_std = target_factors.std(dim=0, unbiased=False)
+            encoder_std = student_embedding.float().std(dim=0, unbiased=False)
+            factor_activity = F.relu(0.02 - factor_std).mean()
+            encoder_variance = F.relu(0.1 - encoder_std).mean()
+        total = prediction + orthogonality + factor_activity + encoder_variance
+        return {
+            "total": total,
+            "prediction": prediction,
+            "orthogonality": orthogonality,
+            "factor_activity": factor_activity,
+            "encoder_variance": encoder_variance,
+        }
+
+    @torch.no_grad()
+    def factor_diagnostics(self) -> dict[str, float]:
+        full = torch.cat(tuple(self.projectors), dim=1).float()
+        gram = full.T @ full
+        singular = torch.linalg.svdvals(full)
+        return {
+            "orthogonality_rmse": float((gram - torch.eye(self.width, device=full.device)).pow(2).mean().sqrt()),
+            "projection_singular_min": float(singular.min()),
+            "projection_singular_max": float(singular.max()),
+            "effective_factors": float(self.num_factors),
+        }

02_Single_Cell/Norman_Single_Cell_Perturbation_Prediction/test/Norman_Single_Cell_Perturbation_Prediction_test.py

@@ -0,0 +1,143 @@
+#!/usr/bin/env python3
+"""Evaluate the archived JEPA-Anything Norman checkpoint.
+
+The archived package contains this script, its model implementation, the
+checkpoint, and the GeneVocab.  GEARS and scGPT remain runtime dependencies in
+the original repository supplied through --repo; the Norman dataset is never
+copied into the archive.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--repo", type=Path, required=True)
+    parser.add_argument("--data-root", type=Path)
+    parser.add_argument("--checkpoint", type=Path, required=True)
+    parser.add_argument("--vocab", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--seed", type=int, default=1)
+    parser.add_argument("--batch-size", type=int, default=64)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--checkpoint-only", action="store_true",
+                        help="Load and validate the frozen model without a dataset")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    archive = Path(__file__).resolve().parents[1]
+    sys.path[:0] = [
+        str(archive / "model"),
+        str(archive / "test"),
+        str(args.repo / "src" / "GEARS"),
+        str(args.repo / "src" / "scGPT"),
+        str(args.repo / "third_party" / "GEARS"),
+        str(args.repo / "third_party" / "scGPT"),
+    ]
+
+    from cell_jepa.objectives_anything import PerturbationJEPAAnything
+    from Norman_Single_Cell_Perturbation_Prediction_test_utils import (
+        add_top20_non_dropout_metrics,
+        build_model,
+        perturb_flags,
+        set_seed,
+    )
+    from scgpt.tokenizer import GeneVocab
+
+    if args.device == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError("--device cuda was requested but CUDA is unavailable")
+    device = torch.device(args.device)
+    torch.set_num_threads(8)
+    torch.set_num_interop_threads(1)
+    set_seed(args.seed)
+
+    payload = torch.load(args.checkpoint, map_location="cpu", weights_only=True)
+    config = payload.get("args", {})
+    num_factors = int(config.get("num_factors", 4))
+    vocab = GeneVocab.from_file(args.vocab)
+    for token in ("<pad>", "<cls>", "<eoc>"):
+        if token not in vocab:
+            vocab.append_token(token)
+    student = build_model(vocab, use_fast=False)
+    model = PerturbationJEPAAnything(student, num_factors=num_factors).to(device)
+    model.load_state_dict(payload["model"], strict=True)
+    model.eval()
+    if args.checkpoint_only:
+        output = {"status": "checkpoint_loaded", "checkpoint": str(args.checkpoint),
+                  "parameters": sum(parameter.numel() for parameter in model.parameters()),
+                  "num_factors": num_factors}
+        args.output.parent.mkdir(parents=True, exist_ok=True)
+        args.output.write_text(json.dumps(output, indent=2) + "\n")
+        print(json.dumps(output, indent=2))
+        return
+    if args.data_root is None:
+        raise ValueError("--data-root is required unless --checkpoint-only is used")
+    from gears import PertData
+    from scgpt.utils import compute_perturbation_metrics, map_raw_id_to_vocab_id
+
+    pert_data = PertData(str(args.data_root))
+    pert_data.default_pert_graph = False
+    pert_data.load(data_path=str(args.data_root / "norman"))
+    pert_data.ctrl_adata = pert_data.adata[pert_data.adata.obs["condition"] == "ctrl"]
+    pert_data.prepare_split(split="simulation", seed=args.seed)
+    pert_data.get_dataloader(batch_size=args.batch_size, test_batch_size=4)
+
+    gene_ids = np.asarray(
+        [vocab[gene] if gene in vocab else vocab["<pad>"]
+         for gene in pert_data.adata.var["gene_name"].tolist()],
+        dtype=int,
+    )
+
+    results: dict[str, list] = {
+        "pert_cat": [], "pred": [], "truth": [], "pred_de": [], "truth_de": [],
+    }
+    with torch.no_grad():
+        for batch in pert_data.dataloader["test_loader"]:
+            batch = batch.to(device)
+            batch_size, n_genes = len(batch.y), batch.y.shape[1]
+            control = batch.x[:, 0].view(batch_size, n_genes)
+            flags = perturb_flags(batch, n_genes, device)
+            raw_ids = torch.arange(n_genes, device=device)
+            genes = map_raw_id_to_vocab_id(raw_ids, gene_ids).repeat(batch_size, 1)
+            padding = torch.zeros_like(control, dtype=torch.bool)
+            prediction = model.predict_expression(genes, control, flags, padding).float()
+            results["pert_cat"].extend(batch.pert)
+            results["pred"].append(prediction.cpu())
+            results["truth"].append(batch.y.cpu())
+            for index, de_idx in enumerate(batch.de_idx):
+                results["pred_de"].append(prediction[index, de_idx].cpu())
+                results["truth_de"].append(batch.y[index, de_idx].cpu())
+
+    results = {
+        key: (np.asarray(value) if key == "pert_cat" else torch.stack(value).numpy()
+              if key in {"pred_de", "truth_de"} else torch.cat(value).numpy())
+        for key, value in results.items()
+    }
+    metrics = compute_perturbation_metrics(results, pert_data.ctrl_adata)
+    metrics = add_top20_non_dropout_metrics(metrics, pert_data.adata, results)
+    output = {
+        "status": "completed",
+        "checkpoint": str(args.checkpoint),
+        "dataset": "norman",
+        "seed": args.seed,
+        "test_cells": int(len(results["pred"])),
+        "test_conditions": int(len(np.unique(results["pert_cat"]))),
+        "metrics": {key: float(value) for key, value in metrics.items()},
+    }
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(json.dumps(output, indent=2) + "\n")
+    print(json.dumps(output, indent=2))
+
+
+if __name__ == "__main__":
+    main()

02_Single_Cell/Norman_Single_Cell_Perturbation_Prediction/test/Norman_Single_Cell_Perturbation_Prediction_test_utils.py

@@ -0,0 +1,67 @@
+"""Evaluation helpers for Norman single-cell perturbation prediction."""
+
+from __future__ import annotations
+
+import random
+
+import numpy as np
+import torch
+from scgpt.model import TransformerGenerator
+from scgpt.tokenizer import GeneVocab
+
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def build_model(vocab: GeneVocab, use_fast: bool) -> TransformerGenerator:
+    return TransformerGenerator(
+        ntoken=len(vocab),
+        d_model=512,
+        nhead=8,
+        d_hid=512,
+        nlayers=12,
+        nlayers_cls=3,
+        n_cls=1,
+        vocab=vocab,
+        dropout=0.0,
+        pad_token="<pad>",
+        pad_value=0,
+        pert_pad_id=0,
+        use_fast_transformer=use_fast,
+    )
+
+
+def perturb_flags(batch, n_genes: int, device: torch.device) -> torch.Tensor:
+    """Support both historic and current GEARS batch layouts."""
+    batch_size = len(batch.y)
+    if batch.x.shape[1] > 1:
+        return batch.x[:, 1].long().view(batch_size, n_genes)
+    flags = torch.zeros((batch_size, n_genes), dtype=torch.long, device=device)
+    for row, indices in enumerate(batch.pert_idx):
+        flags[row, torch.as_tensor(indices, dtype=torch.long, device=device)] = 1
+    return flags
+
+
+def add_top20_non_dropout_metrics(metrics: dict, adata, results: dict) -> dict:
+    """Add GEARS top-20 non-dropout Pearson metrics."""
+    from gears.inference import non_dropout_analysis
+
+    per_perturbation = non_dropout_analysis(adata, results)
+    for output_key, gears_key in (
+        ("top20_de_non_dropout", "pearson_top20_de_non_dropout"),
+        ("pearson_delta_top20_de_non_dropout", "pearson_delta_top20_de_non_dropout"),
+    ):
+        values = [
+            value[gears_key]
+            for value in per_perturbation.values()
+            if gears_key in value
+        ]
+        if values:
+            metrics[output_key] = float(np.mean(values))
+    if "top20_de_non_dropout" in metrics:
+        metrics["top20"] = metrics["top20_de_non_dropout"]
+    return metrics

02_Single_Cell/README.md

@@ -0,0 +1,11 @@
+# Single-Cell Perturbation
+
+Single-cell RNA sequencing describes a cell with a high-dimensional and sparse expression profile. JEPA pretraining predicts a complete cell representation from masked expression, encouraging the encoder to retain biological state while reducing sensitivity to gene-level dropout and count noise. Perturbation experiments use the difference between perturbed and control embeddings as the target response. OPF divides that response into complementary factors that are recombined to predict the post-perturbation expression profile.
+
+Open a model page for its research background, checkpoint, input format, and inference command. Additional checkpoints from this domain will be released progressively.
+
+## Models
+
+- [Norman_Single_Cell_Perturbation_Prediction](Norman_Single_Cell_Perturbation_Prediction/README.md)
+
+[Back to the project overview](../README.md)

04_CITRIS_Intervention/CITRIS_Interventional_Pong/README.md

@@ -0,0 +1,36 @@
+# CITRIS Interventional Pong — Legacy Model
+
+## Background
+
+Interventional Pong exposes interventions on ball motion, ball position, and paddle position. The task studies whether single-intervention dynamics can be recombined when several variables change together. This checkpoint belongs to the earlier three-channel protocol.
+
+## Files
+
+- Model implementation: [`../../model`](<../../model>)
+- Checkpoint: [`checkpoints/CITRIS_Interventional_Pong_final_checkpoint.pt`](checkpoints/CITRIS_Interventional_Pong_final_checkpoint.pt)
+- Inference command: [`../../test/infer_dynamics.py`](<../../test/infer_dynamics.py>)
+
+## Run inference
+
+Run from this experiment directory:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/CITRIS_Interventional_Pong_final_checkpoint.pt \
+  --input <preprocessed_transitions.npz> \
+  --output <new_output_dir>/predictions.npz
+```
+
+Provide `x` and `action` for one-step prediction, or `x` and `actions` for a recurrent rollout. Optional `y` or `targets` arrays add overall and per-step MSE to the output report. Use the dataset-specific flattening, channel order, normalization, and action scaling described for the experiment.
+
+For a quick checkpoint and tensor-shape check, omit `--input`:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/CITRIS_Interventional_Pong_final_checkpoint.pt \
+  --output /tmp/jepa_inference_smoke/predictions.npz
+```
+
+This command reports the predicted tensor shape and finite-output status. Dataset runs also report MSE when targets are supplied.
+
+[Back to the project overview](<../../README.md>)

04_CITRIS_Intervention/CITRIS_Interventional_Pong/checkpoints/CITRIS_Interventional_Pong_final_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:80c890f601d4af364850737298307312f42203cef9fd2c0baaee5cb782feeb54
+size 14709115

04_CITRIS_Intervention/README.md

@@ -0,0 +1,12 @@
+# Interventional Dynamics
+
+CITRIS Interventional Pong separates ordinary temporal evolution from explicit interventions on underlying game variables such as ball motion, ball position, and paddle position. The model receives the pre-intervention observation together with a label describing which variables changed, then predicts the resulting next state. This makes the task a direct test of compositional dynamics: factors learned from individual interventions should remain useful when several interventions occur together.
+
+Open a model page for its research background, checkpoint, input format, and inference command. Additional checkpoints from this domain will be released progressively.
+
+## Models
+
+- [CITRIS_Interventional_Pong](CITRIS_Interventional_Pong/README.md)
+- [six_step_one_step](six_step_one_step/README.md)
+
+[Back to the project overview](../README.md)

04_CITRIS_Intervention/six_step_one_step/README.md

@@ -0,0 +1,35 @@
+# CITRIS Six-Step Models: Strict One-Step Evaluation
+
+## Background
+
+These models were trained with recurrent six-step feedback and are evaluated here on one-step prediction. The experiment asks whether a model trained for short rollout can predict held-out single interventions and simultaneous intervention combinations.
+
+The protocol uses all four original channels: RGB plus the movement-projection channel. It is separate from the legacy three-channel checkpoint.
+
+## Files
+
+- Inference model: [`../../model`](../../model)
+- Checkpoints: [`checkpoints`](checkpoints), with Dense and OPF weights for seeds 11, 23, 37, 53, and 71
+- Test entry point: [`test/evaluate.py`](test/evaluate.py)
+
+## Run the test
+
+```bash
+python test/evaluate.py \
+  --data-root <citris_data_directory> \
+  --output <new_output_directory> \
+  --device cuda
+```
+
+Provide the original four-channel `single_test.npz` and `combo_test.npz` files. The script loads ten checkpoints, evaluates exactly-one-intervention transitions and transitions with at least two interventions, and writes `results.json` plus `per_seed.csv`.
+
+Five-seed four-channel MSE from the verified checkpoint run:
+
+| One-step test split | Dense | OPF |
+|---|---:|---:|
+| Exactly one intervention | 0.00952107 | 0.00682598 |
+| At least two interventions | 0.00942106 | 0.00822350 |
+
+The table reports one-step evaluation for models trained with recurrent six-step feedback.
+
+[Back to the project overview](../../README.md)

04_CITRIS_Intervention/six_step_one_step/checkpoints/seed11/h6_dense.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3f688358c41f5f3b2ed3e78c844c0e3883b421d083e6ff40d809a312c28dcad7
+size 18642739

04_CITRIS_Intervention/six_step_one_step/checkpoints/seed11/h6_opf.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bbbe28b515c5152b679c4e763bddb47bac58639324d366e5c9257dfee2a4e1bd
+size 18640945

04_CITRIS_Intervention/six_step_one_step/checkpoints/seed23/h6_dense.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1bf5b89faacf28cf99ab82b85f8f14a31276eeeb01d173aaee69e42fde9085d5
+size 18642739

04_CITRIS_Intervention/six_step_one_step/checkpoints/seed23/h6_opf.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3d8138d3c13fd448e7a2986cd1fcce18fceeb3c31362f3e7a897e0771217f105
+size 18640945

04_CITRIS_Intervention/six_step_one_step/checkpoints/seed37/h6_dense.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fa880b7b2e6610670da9b7b49f98dd24b2100dd98091d19d78bac5add140a688
+size 18642739

04_CITRIS_Intervention/six_step_one_step/checkpoints/seed37/h6_opf.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1170a0b22b2048b761cb07b56175d2508a8c0c27fc57277ac040e9fac1ea2f88
+size 18640945

04_CITRIS_Intervention/six_step_one_step/checkpoints/seed53/h6_dense.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe47c9b5826b5d6547d324b736dcff46894152b45c4a262b229b5e37077882d1
+size 18642739

04_CITRIS_Intervention/six_step_one_step/checkpoints/seed53/h6_opf.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:862a2d7db70b4ed8992ae681c26815216d0d687998198eff13f01b4308a15c1f
+size 18640945

04_CITRIS_Intervention/six_step_one_step/checkpoints/seed71/h6_dense.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6e0953e2d690607efd5c04c58903d42be2e4643448e50c25de6c1332505d2f7a
+size 18642739

04_CITRIS_Intervention/six_step_one_step/checkpoints/seed71/h6_opf.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:86e1904460c57e6360951d62277431ab3273500e3493225cc91dca3a173793f6
+size 18640945

04_CITRIS_Intervention/six_step_one_step/test/evaluate.py

@@ -0,0 +1,60 @@
+#!/usr/bin/env python3
+"""Evaluate the frozen six-step CITRIS checkpoints on strict one-step splits. No training."""
+from __future__ import annotations
+import argparse,csv,json,sys
+from pathlib import Path
+import numpy as np
+import torch
+
+PROJECT=Path(__file__).resolve().parents[3]
+sys.path.insert(0,str(PROJECT))
+from model import load_dynamics_checkpoint
+
+GROUPS=['ball_vel_dir','ball_x','ball_y','paddle_left_y','paddle_right_y']
+SEEDS=[11,23,37,53,71]
+
+def load_split(path):
+    with np.load(path,allow_pickle=False) as data:
+        ids=[data['keys'].tolist().index(k) for k in GROUPS]
+        x=torch.from_numpy(data['images'].copy()).permute(0,3,1,2).float().div_(255).flatten(1)
+        actions=torch.from_numpy(data['targets'][:,ids].astype('float32'))
+    if x.shape[1]!=4096 or len(actions)!=len(x)-1:raise ValueError('Expected original four-channel CITRIS arrays')
+    return x,actions
+
+def mse(model,x,actions,indices,device,batch=256):
+    total=0.0
+    with torch.inference_mode():
+        for start in range(0,len(indices),batch):
+            ix=indices[start:start+batch]
+            prediction=model(x[ix].to(device),actions[ix].to(device))
+            total+=float((prediction-x[ix+1].to(device)).square().sum())
+    return total/(len(indices)*x.shape[1])
+
+def main():
+    p=argparse.ArgumentParser(description=__doc__);p.add_argument('--data-root',type=Path,required=True);p.add_argument('--checkpoint-root',type=Path,default=Path(__file__).resolve().parents[1]/'checkpoints');p.add_argument('--output',type=Path,required=True);p.add_argument('--device',default='cpu');a=p.parse_args()
+    rows=[]
+    splits=[('single_test',a.data_root/'single_test.npz',lambda active:active==1),('combo_test',a.data_root/'combo_test.npz',lambda active:active>=2)]
+    loaded_splits={}
+    for name,path,rule in splits:
+        x,actions=load_split(path);loaded_splits[name]=(x,actions,torch.where(rule(actions.sum(1)))[0])
+    for seed in SEEDS:
+        for kind in ['dense','opf']:
+            path=a.checkpoint_root/f'seed{seed}'/f'h6_{kind}.pt'
+            loaded=load_dynamics_checkpoint(path,'model',a.device)
+            if loaded.metadata.get('horizon')!=6 or loaded.metadata.get('kind')!=kind:raise ValueError(f'Unexpected metadata in {path}')
+            for split,(x,actions,indices) in loaded_splits.items():
+                rows.append({'seed':seed,'model':kind,'split':split,'count':len(indices),'mse_all4':mse(loaded.model,x,actions,indices,a.device)})
+    summary={}
+    for split,_,_ in splits:
+        summary[split]={}
+        for kind in ['dense','opf']:
+            values=np.array([row['mse_all4'] for row in rows if row['split']==split and row['model']==kind])
+            summary[split][kind]={'mean':float(values.mean()),'sample_sd':float(values.std(ddof=1)),'values':values.tolist()}
+        summary[split]['opf_reduction_percent']=float(100*(1-summary[split]['opf']['mean']/summary[split]['dense']['mean']))
+    a.output.mkdir(parents=True,exist_ok=False)
+    with (a.output/'per_seed.csv').open('w',newline='') as f:
+        writer=csv.DictWriter(f,fieldnames=list(rows[0]));writer.writeheader();writer.writerows(rows)
+    (a.output/'results.json').write_text(json.dumps({'status':'complete','training_horizon':6,'evaluation_horizon':1,'summary':summary},indent=2))
+    print(json.dumps(summary,indent=2))
+
+if __name__=='__main__':main()

05_Ten_Task_Dynamics/CLEVRER_Visual_Dynamics/README.md

@@ -0,0 +1,36 @@
+# CLEVRER Visual Dynamics
+
+## Background
+
+CLEVRER represents dynamics as controlled video transitions. The checkpoint predicts the next visual state and can be rolled forward to study error accumulation.
+
+## Files
+
+- Model implementation: [`../../model`](<../../model>)
+- Checkpoint: [`checkpoints/CLEVRER_Visual_Dynamics_final_checkpoint.pt`](checkpoints/CLEVRER_Visual_Dynamics_final_checkpoint.pt)
+- Inference command: [`../../test/infer_dynamics.py`](<../../test/infer_dynamics.py>)
+
+## Run inference
+
+Run from this experiment directory:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/CLEVRER_Visual_Dynamics_final_checkpoint.pt \
+  --input <preprocessed_transitions.npz> \
+  --output <new_output_dir>/predictions.npz
+```
+
+Provide `x` and `action` for one-step prediction, or `x` and `actions` for a recurrent rollout. Optional `y` or `targets` arrays add overall and per-step MSE to the output report. Use the dataset-specific flattening, channel order, normalization, and action scaling described for the experiment.
+
+For a quick checkpoint and tensor-shape check, omit `--input`:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/CLEVRER_Visual_Dynamics_final_checkpoint.pt \
+  --output /tmp/jepa_inference_smoke/predictions.npz
+```
+
+This command reports the predicted tensor shape and finite-output status. Dataset runs also report MSE when targets are supplied.
+
+[Back to the project overview](<../../README.md>)

05_Ten_Task_Dynamics/CLEVRER_Visual_Dynamics/checkpoints/CLEVRER_Visual_Dynamics_final_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8251b743f84ba43c8ab94307a9f44b6242970f36df1ffbf27268cc031c5d9e9d
+size 39657359

05_Ten_Task_Dynamics/CausalWorld_Pushing/README.md

@@ -0,0 +1,36 @@
+# CausalWorld Pushing
+
+## Background
+
+This task predicts how a robot action changes object state. It tests whether action-conditioned latent factors retain the physical motion signal across controlled conditions.
+
+## Files
+
+- Model implementation: [`../../model`](<../../model>)
+- Checkpoint: [`checkpoints/CausalWorld_Pushing_final_checkpoint.pt`](checkpoints/CausalWorld_Pushing_final_checkpoint.pt)
+- Inference command: [`../../test/infer_dynamics.py`](<../../test/infer_dynamics.py>)
+
+## Run inference
+
+Run from this experiment directory:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/CausalWorld_Pushing_final_checkpoint.pt \
+  --input <preprocessed_transitions.npz> \
+  --output <new_output_dir>/predictions.npz
+```
+
+Provide `x` and `action` for one-step prediction, or `x` and `actions` for a recurrent rollout. Optional `y` or `targets` arrays add overall and per-step MSE to the output report. Use the dataset-specific flattening, channel order, normalization, and action scaling described for the experiment.
+
+For a quick checkpoint and tensor-shape check, omit `--input`:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/CausalWorld_Pushing_final_checkpoint.pt \
+  --output /tmp/jepa_inference_smoke/predictions.npz
+```
+
+This command reports the predicted tensor shape and finite-output status. Dataset runs also report MSE when targets are supplied.
+
+[Back to the project overview](<../../README.md>)

05_Ten_Task_Dynamics/CausalWorld_Pushing/checkpoints/CausalWorld_Pushing_final_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd7b198d91465e6284cddafbfa89ecc1d9cf2ede1678eeb89714f485b61b4d35
+size 2039887

05_Ten_Task_Dynamics/CausalWorld_Pushing_Long_Horizon/README.md

@@ -0,0 +1,36 @@
+# CausalWorld Long-Horizon Pushing
+
+## Background
+
+This variant repeatedly feeds model predictions back into the dynamics model. It is intended for studying long-horizon state prediction and model-based control.
+
+## Files
+
+- Model implementation: [`../../model`](<../../model>)
+- Checkpoint: [`checkpoints/CausalWorld_Pushing_Long_Horizon_final_checkpoint.pt`](checkpoints/CausalWorld_Pushing_Long_Horizon_final_checkpoint.pt)
+- Inference command: [`../../test/infer_dynamics.py`](<../../test/infer_dynamics.py>)
+
+## Run inference
+
+Run from this experiment directory:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/CausalWorld_Pushing_Long_Horizon_final_checkpoint.pt \
+  --input <preprocessed_transitions.npz> \
+  --output <new_output_dir>/predictions.npz
+```
+
+Provide `x` and `action` for one-step prediction, or `x` and `actions` for a recurrent rollout. Optional `y` or `targets` arrays add overall and per-step MSE to the output report. Use the dataset-specific flattening, channel order, normalization, and action scaling described for the experiment.
+
+For a quick checkpoint and tensor-shape check, omit `--input`:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/CausalWorld_Pushing_Long_Horizon_final_checkpoint.pt \
+  --output /tmp/jepa_inference_smoke/predictions.npz
+```
+
+This command reports the predicted tensor shape and finite-output status. Dataset runs also report MSE when targets are supplied.
+
+[Back to the project overview](<../../README.md>)

05_Ten_Task_Dynamics/CausalWorld_Pushing_Long_Horizon/checkpoints/CausalWorld_Pushing_Long_Horizon_final_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd7f2a595e2feae5b5827ee7ee6ff4c07c7188fe0bffaf9a398493cf67cd94c3
+size 2039951

05_Ten_Task_Dynamics/DMC_Cartpole_Pixels/README.md

@@ -0,0 +1,36 @@
+# DMC Cartpole from Pixels
+
+## Background
+
+The cartpole state is observed as pixels rather than low-dimensional coordinates. The model predicts how an action changes the next visual state.
+
+## Files
+
+- Model implementation: [`../../model`](<../../model>)
+- Checkpoint: [`checkpoints/DMC_Cartpole_Pixels_final_checkpoint.pt`](checkpoints/DMC_Cartpole_Pixels_final_checkpoint.pt)
+- Inference command: [`../../test/infer_dynamics.py`](<../../test/infer_dynamics.py>)
+
+## Run inference
+
+Run from this experiment directory:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/DMC_Cartpole_Pixels_final_checkpoint.pt \
+  --input <preprocessed_transitions.npz> \
+  --output <new_output_dir>/predictions.npz
+```
+
+Provide `x` and `action` for one-step prediction, or `x` and `actions` for a recurrent rollout. Optional `y` or `targets` arrays add overall and per-step MSE to the output report. Use the dataset-specific flattening, channel order, normalization, and action scaling described for the experiment.
+
+For a quick checkpoint and tensor-shape check, omit `--input`:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/DMC_Cartpole_Pixels_final_checkpoint.pt \
+  --output /tmp/jepa_inference_smoke/predictions.npz
+```
+
+This command reports the predicted tensor shape and finite-output status. Dataset runs also report MSE when targets are supplied.
+
+[Back to the project overview](<../../README.md>)

05_Ten_Task_Dynamics/DMC_Cartpole_Pixels/checkpoints/DMC_Cartpole_Pixels_final_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0d0ec41d94f43125aadf995efea8f73a9ec0173a7e0b1afb86e51ed0cda49d12
+size 39657359

05_Ten_Task_Dynamics/DMC_VB_Cheetah_Run/README.md

@@ -0,0 +1,36 @@
+# DMC-VB Cheetah Run
+
+## Background
+
+This task combines locomotion dynamics with changing visual backgrounds. It studies whether the predictor follows body motion instead of appearance alone.
+
+## Files
+
+- Model implementation: [`../../model`](<../../model>)
+- Checkpoint: [`checkpoints/DMC_VB_Cheetah_Run_final_checkpoint.pt`](checkpoints/DMC_VB_Cheetah_Run_final_checkpoint.pt)
+- Inference command: [`../../test/infer_dynamics.py`](<../../test/infer_dynamics.py>)
+
+## Run inference
+
+Run from this experiment directory:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/DMC_VB_Cheetah_Run_final_checkpoint.pt \
+  --input <preprocessed_transitions.npz> \
+  --output <new_output_dir>/predictions.npz
+```
+
+Provide `x` and `action` for one-step prediction, or `x` and `actions` for a recurrent rollout. Optional `y` or `targets` arrays add overall and per-step MSE to the output report. Use the dataset-specific flattening, channel order, normalization, and action scaling described for the experiment.
+
+For a quick checkpoint and tensor-shape check, omit `--input`:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/DMC_VB_Cheetah_Run_final_checkpoint.pt \
+  --output /tmp/jepa_inference_smoke/predictions.npz
+```
+
+This command reports the predicted tensor shape and finite-output status. Dataset runs also report MSE when targets are supplied.
+
+[Back to the project overview](<../../README.md>)

05_Ten_Task_Dynamics/DMC_VB_Cheetah_Run/checkpoints/DMC_VB_Cheetah_Run_final_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:62fc08813c234b473bb5ab4b535e0fa6e81fe6b7f9302540d539c85ba7d1baa1
+size 39662479

05_Ten_Task_Dynamics/Distracting_Control_Cartpole/README.md

@@ -0,0 +1,36 @@
+# Distracting Control Cartpole
+
+## Background
+
+The physical cartpole evolves while backgrounds and visual distractors vary. The model is used to separate predictable motion from nuisance appearance.
+
+## Files
+
+- Model implementation: [`../../model`](<../../model>)
+- Checkpoint: [`checkpoints/Distracting_Control_Cartpole_final_checkpoint.pt`](checkpoints/Distracting_Control_Cartpole_final_checkpoint.pt)
+- Inference command: [`../../test/infer_dynamics.py`](<../../test/infer_dynamics.py>)
+
+## Run inference
+
+Run from this experiment directory:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/Distracting_Control_Cartpole_final_checkpoint.pt \
+  --input <preprocessed_transitions.npz> \
+  --output <new_output_dir>/predictions.npz
+```
+
+Provide `x` and `action` for one-step prediction, or `x` and `actions` for a recurrent rollout. Optional `y` or `targets` arrays add overall and per-step MSE to the output report. Use the dataset-specific flattening, channel order, normalization, and action scaling described for the experiment.
+
+For a quick checkpoint and tensor-shape check, omit `--input`:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/Distracting_Control_Cartpole_final_checkpoint.pt \
+  --output /tmp/jepa_inference_smoke/predictions.npz
+```
+
+This command reports the predicted tensor shape and finite-output status. Dataset runs also report MSE when targets are supplied.
+
+[Back to the project overview](<../../README.md>)

05_Ten_Task_Dynamics/Distracting_Control_Cartpole/checkpoints/Distracting_Control_Cartpole_final_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d19c4126fb73ac9bc43e7ee43c6990a33bbc4f927c9f29a2e362998e3e7da4c8
+size 39657423

05_Ten_Task_Dynamics/PDEBench_Burgers/README.md

@@ -0,0 +1,36 @@
+# PDEBench Burgers
+
+## Background
+
+Burgers dynamics combine nonlinear transport and dissipation. The checkpoint predicts the next discretized field and supports recurrent rollout.
+
+## Files
+
+- Model implementation: [`../../model`](<../../model>)
+- Checkpoint: [`checkpoints/PDEBench_Burgers_final_checkpoint.pt`](checkpoints/PDEBench_Burgers_final_checkpoint.pt)
+- Inference command: [`../../test/infer_dynamics.py`](<../../test/infer_dynamics.py>)
+
+## Run inference
+
+Run from this experiment directory:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/PDEBench_Burgers_final_checkpoint.pt \
+  --input <preprocessed_transitions.npz> \
+  --output <new_output_dir>/predictions.npz
+```
+
+Provide `x` and `action` for one-step prediction, or `x` and `actions` for a recurrent rollout. Optional `y` or `targets` arrays add overall and per-step MSE to the output report. Use the dataset-specific flattening, channel order, normalization, and action scaling described for the experiment.
+
+For a quick checkpoint and tensor-shape check, omit `--input`:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/PDEBench_Burgers_final_checkpoint.pt \
+  --output /tmp/jepa_inference_smoke/predictions.npz
+```
+
+This command reports the predicted tensor shape and finite-output status. Dataset runs also report MSE when targets are supplied.
+
+[Back to the project overview](<../../README.md>)

05_Ten_Task_Dynamics/PDEBench_Burgers/checkpoints/PDEBench_Burgers_final_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:83bbc2cacc9ff7769521a76c3c973c4b3895dd92362a49f2700cd80f8b334826
+size 5009359

05_Ten_Task_Dynamics/PDEBench_Navier_Stokes/README.md

@@ -0,0 +1,36 @@
+# PDEBench Navier–Stokes
+
+## Background
+
+This task predicts the evolution of a two-dimensional fluid field. It tests whether the latent state retains spatial structure during repeated prediction.
+
+## Files
+
+- Model implementation: [`../../model`](<../../model>)
+- Checkpoint: [`checkpoints/PDEBench_Navier_Stokes_final_checkpoint.pt`](checkpoints/PDEBench_Navier_Stokes_final_checkpoint.pt)
+- Inference command: [`../../test/infer_dynamics.py`](<../../test/infer_dynamics.py>)
+
+## Run inference
+
+Run from this experiment directory:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/PDEBench_Navier_Stokes_final_checkpoint.pt \
+  --input <preprocessed_transitions.npz> \
+  --output <new_output_dir>/predictions.npz
+```
+
+Provide `x` and `action` for one-step prediction, or `x` and `actions` for a recurrent rollout. Optional `y` or `targets` arrays add overall and per-step MSE to the output report. Use the dataset-specific flattening, channel order, normalization, and action scaling described for the experiment.
+
+For a quick checkpoint and tensor-shape check, omit `--input`:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/PDEBench_Navier_Stokes_final_checkpoint.pt \
+  --output /tmp/jepa_inference_smoke/predictions.npz
+```
+
+This command reports the predicted tensor shape and finite-output status. Dataset runs also report MSE when targets are supplied.
+
+[Back to the project overview](<../../README.md>)

05_Ten_Task_Dynamics/PDEBench_Navier_Stokes/checkpoints/PDEBench_Navier_Stokes_final_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c63e159cef39d2611aec13a1c967d5f506cd3650ea537cffb062bf8f95ca0b18
+size 8159247

05_Ten_Task_Dynamics/PDEBench_Reaction_Diffusion_1D/README.md

@@ -0,0 +1,36 @@
+# PDEBench Reaction–Diffusion 1D
+
+## Background
+
+Reaction and spatial diffusion operate at different scales. The model predicts the next one-dimensional field from the current field.
+
+## Files
+
+- Model implementation: [`../../model`](<../../model>)
+- Checkpoint: [`checkpoints/PDEBench_Reaction_Diffusion_1D_final_checkpoint.pt`](checkpoints/PDEBench_Reaction_Diffusion_1D_final_checkpoint.pt)
+- Inference command: [`../../test/infer_dynamics.py`](<../../test/infer_dynamics.py>)
+
+## Run inference
+
+Run from this experiment directory:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/PDEBench_Reaction_Diffusion_1D_final_checkpoint.pt \
+  --input <preprocessed_transitions.npz> \
+  --output <new_output_dir>/predictions.npz
+```
+
+Provide `x` and `action` for one-step prediction, or `x` and `actions` for a recurrent rollout. Optional `y` or `targets` arrays add overall and per-step MSE to the output report. Use the dataset-specific flattening, channel order, normalization, and action scaling described for the experiment.
+
+For a quick checkpoint and tensor-shape check, omit `--input`:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/PDEBench_Reaction_Diffusion_1D_final_checkpoint.pt \
+  --output /tmp/jepa_inference_smoke/predictions.npz
+```
+
+This command reports the predicted tensor shape and finite-output status. Dataset runs also report MSE when targets are supplied.
+
+[Back to the project overview](<../../README.md>)

05_Ten_Task_Dynamics/PDEBench_Reaction_Diffusion_1D/checkpoints/PDEBench_Reaction_Diffusion_1D_final_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7db8dbc4ce2f77f525b34a2b69bb03754aa27d6358e3a8f1a1bd6a704db81535
+size 5009359

05_Ten_Task_Dynamics/PDEBench_Reaction_Diffusion_2D/README.md

@@ -0,0 +1,36 @@
+# PDEBench Reaction–Diffusion 2D
+
+## Background
+
+This task extends reaction–diffusion forecasting to two-dimensional spatial fields and recurrent prediction.
+
+## Files
+
+- Model implementation: [`../../model`](<../../model>)
+- Checkpoint: [`checkpoints/PDEBench_Reaction_Diffusion_2D_final_checkpoint.pt`](checkpoints/PDEBench_Reaction_Diffusion_2D_final_checkpoint.pt)
+- Inference command: [`../../test/infer_dynamics.py`](<../../test/infer_dynamics.py>)
+
+## Run inference
+
+Run from this experiment directory:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/PDEBench_Reaction_Diffusion_2D_final_checkpoint.pt \
+  --input <preprocessed_transitions.npz> \
+  --output <new_output_dir>/predictions.npz
+```
+
+Provide `x` and `action` for one-step prediction, or `x` and `actions` for a recurrent rollout. Optional `y` or `targets` arrays add overall and per-step MSE to the output report. Use the dataset-specific flattening, channel order, normalization, and action scaling described for the experiment.
+
+For a quick checkpoint and tensor-shape check, omit `--input`:
+
+```bash
+python ../../test/infer_dynamics.py \
+  --checkpoint checkpoints/PDEBench_Reaction_Diffusion_2D_final_checkpoint.pt \
+  --output /tmp/jepa_inference_smoke/predictions.npz
+```
+
+This command reports the predicted tensor shape and finite-output status. Dataset runs also report MSE when targets are supplied.
+
+[Back to the project overview](<../../README.md>)

05_Ten_Task_Dynamics/PDEBench_Reaction_Diffusion_2D/checkpoints/PDEBench_Reaction_Diffusion_2D_final_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bcaba73788e20dcae2f273a803765d1f6abec6e280d8064ab3c4a159e9d25f9e
+size 8159183