aguennoune17/atlas-v2-nwm-fp8-compressed

Baseline ATLAS v2:  31.7%
After Phase 1:      50.0% (+18.3%)
After Phase 2:      65.0% (+15.0%)

Input: context_emb (64D)
├── h1 branch: Linear(64→128) + LayerNorm + GELU + Linear(128→64) → pred_h1
├── h2 branch: Linear(64→128) + LayerNorm + GELU + Linear(128→64) → pred_h2
└── h4 branch: Linear(64→128) + LayerNorm + GELU + Linear(128→64) → pred_h4

{
  'latent_dim': 32,
  'mask_ratio': 0.4,
  'block_size': 4,
  'batch_size': 128,
  'num_steps': 200,
  'lr_max': 1e-4,
  'lr_min': 1e-5,
  'warmup_steps': 20,
  'ema_momentum': 0.996,
  'lambda_var': 25.0,
  'lambda_cov': 1.0,
  'lambda_vicreg': 0.1,
}

{
  'latent_dim': 64,
  'horizons': [1, 2, 4],
  'horizon_weights': [1.0, 0.8, 0.5],
  'mask_ratio': 0.4,
  'block_size': 8,
  'batch_size': 64,
  'num_steps': 500,
  'lr_max': 1e-4,
  'lr_min': 1e-5,
  'warmup_steps': 50,
  'ema_momentum': 0.996,
  'lambda_var': 25.0,
  'lambda_cov': 1.0,
  'lambda_vicreg': 0.1,
}

  loss_mh = w1 × MSE(pred_h1, target_h1) + 
            w2 × MSE(pred_h2, target_h2) +
            w4 × MSE(pred_h4, target_h4)

pip install torch torchvision einops

import torch

# Load checkpoint
checkpoint = torch.load('atlas_phase1_checkpoint.pt', map_location='cpu')

# Extract components
context_encoder_state = checkpoint['context_encoder']
target_encoder_state = checkpoint['target_encoder']
config = checkpoint['config']
history = checkpoint['history']

print(f"Phase 1 — Latent dim: {config['latent_dim']}")
print(f"Final loss: {history['loss_total'][-1]:.4f}")

import torch

# Load checkpoint
checkpoint = torch.load('atlas_phase2_checkpoint.pt', map_location='cpu')

# Extract components
context_encoder_state = checkpoint['context_encoder']
target_encoder_state = checkpoint['target_encoder']
multi_horizon_predictor_state = checkpoint['multi_horizon_predictor']
config = checkpoint['config']
history = checkpoint['history']

print(f"Phase 2 — Latent dim: {config['latent_dim']}")
print(f"Horizons: {config['horizons']}")
print(f"Final loss: {history['loss_total'][-1]:.4f}")

import torch
import torch.nn as nn

# Define ContextEncoder (Phase 2)
class ContextEncoder(nn.Module):
    def __init__(self, input_dim=3, latent_dim=64, hidden=128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden),
            nn.LayerNorm(hidden),
            nn.GELU(),
            nn.Linear(hidden, latent_dim),
            nn.LayerNorm(latent_dim),
        )
    
    def forward(self, x):
        return self.net(x)

# Load Phase 2 model
checkpoint = torch.load('atlas_phase2_checkpoint.pt', map_location='cpu')
context_encoder = ContextEncoder(latent_dim=64)
context_encoder.load_state_dict(checkpoint['context_encoder'])
context_encoder.eval()

# Inference
candidate = torch.tensor([[0.75, 0.82, 0.28]])  # [Stability, Alignment, Cost]
with torch.no_grad():
    embedding = context_encoder(candidate)
    print(f"Embedding shape: {embedding.shape}")  # (1, 64)
    print(f"Embedding norm: {embedding.norm().item():.4f}")

# Define MultiHorizonPredictor
class MultiHorizonPredictor(nn.Module):
    def __init__(self, latent_dim=64, hidden=128, horizons=[1, 2, 4]):
        super().__init__()
        self.horizons = horizons
        self.predictors = nn.ModuleDict({
            f'h{h}': nn.Sequential(
                nn.Linear(latent_dim, hidden),
                nn.LayerNorm(hidden),
                nn.GELU(),
                nn.Linear(hidden, latent_dim),
            )
            for h in horizons
        })
    
    def forward(self, context_emb):
        return {h: self.predictors[f'h{h}'](context_emb) for h in self.horizons}

# Load predictor
predictor = MultiHorizonPredictor(latent_dim=64)
predictor.load_state_dict(checkpoint['multi_horizon_predictor'])
predictor.eval()

# Predict future states
with torch.no_grad():
    context_emb = context_encoder(candidate)
    predictions = predictor(context_emb)
    
    print("Predicted embeddings:")
    for h, pred_emb in predictions.items():
        print(f"  t+{h}: {pred_emb.shape}, norm={pred_emb.norm().item():.4f}")

@software{atlas_lewm_2026,
  title={ATLAS v2 + Le WorldModel: Fine-Tuned JEPA Architecture},
  author={Guennoune, Abderrahim and Contributors},
  year={2026},
  month={March},
  publisher={HuggingFace},
  doi={10.5281/zenodo.17764881},
  url={https://huggingface.co/aguennoune/atlas-v2-lewm}
}

@article{lecun2026leworldmodel,
  title={A Clean Recipe to Train JEPA World Models},
  author={LeCun, Yann and others},
  journal={arXiv preprint arXiv:2603.19312},
  year={2026},
  url={https://le-wm.github.io/}
}