1. Overview

This document outlines the architecture for the Omega-JEPA Predictor Network, a clean-room implementation inspired by Joint Embedding Predictive Architectures (JEPA) but specialized for the "Dream Stream" environment. The core innovation is the integration of Protocol Omega metrics (Time-Reversal Asymmetry and Z-Scores) directly into the validation loop to ensure causal fidelity in state predictions.

2. Theoretical Foundation

2.1. The JEPA Paradigm

Unlike generative models that predict pixels (x), JEPA predicts representations (y) in an abstract space.

Context ($S_x$): The current state representation.
Action ($a$): The control or transition vector.
Latent ($z$): A stochastic variable capturing uncertainty in the transition.
Prediction ($S_y$): The predicted future state representation.

Equation: $S_y = Predictor(S_x, a, z)$

2.2. Protocol Omega Integration

To prevent the model from learning "easy" but non-causal shortcuts (representation collapse), we enforce Time-Reversal Asymmetry (TRA).

TRA Hypothesis: A valid causal transition $A \to B$ should be distinguishable from $B \to A$ in the latent energy landscape.
Omega Score: A composite metric combining prediction error (L2) with TRA violation penalties.

3. Architecture Design

3.1. `OmegaJEPA_Predictor` (The Brain)

Type: PyTorch nn.Module.
Structure: Multi-Layer Perceptron (MLP) with residual connections.
Inputs:
- context_embedding: Tensor [Batch, Dim]
- action_vector: Tensor [Batch, ActionDim]
- latent_z: Tensor [Batch, LatentDim] (Optional/sampled)
Mechanism:
1. Concatenate $S_x$, $a$, and $z$.
2. Pass through a high-dimensional projection layer.
3. Apply LayerNorm and GeLU activations.
4. Output projected state $\hat{S}_y$.

3.2. `OmegaMetrics` (The Auditor)

Purpose: Stateless validator class to compute physics-inspired metrics.
Key Metrics:
- compute_tra(state_t, state_t1): Measures asymmetry magnitude.
- compute_z_score(residuals): Standard deviation based anomaly detection.
- energy_function(state): Helper for TRA computation (e.g., magnitude or learned energy).

4. Implementation Constraints

Framework: PyTorch (v2.x).
Style: Strict typing, modular design.
Clean Room: No usage of Meta/Facebook source code; strictly first-principles implementation based on LeCun's 2022 paper and Protocol Omega specs.

5. Usage

model = OmegaJEPA_Predictor(dim=256, action_dim=64)
metrics = OmegaMetrics()

# Forward pass
s_next_pred = model(s_curr, action, z)

# Validation
tra_score = metrics.compute_tra(s_curr, s_next_pred)