"""Diagnose full training path including ACT wrapper and loss head."""

import torch

import torch.nn as nn

from models.recursive_reasoning.trm_vjepa_post_nl import (

VJEPARecursiveConfig,

VJEPARecursiveClassifier_ACTV1,

)

from models.losses import VideoClassificationACTLossHead

def main():

device = "cuda"

torch.manual_seed(42)

# Create model with same config as training

config_dict = dict(

batch_size=4,

hidden_size=512,

vjepa_hidden_size=1024,

num_heads=8,

H_cycles=2,

G_cycles=2,

L_cycles=2,

L_layers=2,

G_level=False,

num_classes=174,

halt_max_steps=8,

halt_exploration_prob=0.2,

)

# Create model with loss head (same as training)

model = VJEPARecursiveClassifier_ACTV1(config_dict).to(device)

model_with_loss = VideoClassificationACTLossHead(model, label_smoothing=0.1, q_loss_weight=0.5)

model_with_loss.train()

print("Model created successfully")

# Fake V-JEPA features and labels - just ONE batch to overfit

B, S, D = 8, 2048, 1024

vjepa_features = torch.randn(B, S, D, device=device, dtype=torch.bfloat16)

labels = torch.randint(0, 174, (B,), device=device)

print(f"Created 1 batch of {B} samples to overfit")

# Test WITHOUT ACT first (simpler, should definitely work)

print("\n=== Test without ACT (fresh carry each step) ===")

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for step in range(50):

# Fresh carry each step

carry = model_with_loss.initial_carry(B, S, vjepa_features, labels)

optimizer.zero_grad()

carry, loss, metrics, outputs, all_halted = model_with_loss(

return_keys=["preds"],

carry=carry,

vjepa_features=vjepa_features,

labels=labels,

)

(loss / B).backward()

optimizer.step()

# Compute accuracy directly from predictions

preds = outputs.get("preds")

if preds is not None:

acc = (preds == labels).float().mean().item()

else:

acc = 0.0

if step % 10 == 0:

print(f"Step {step}: loss={loss.item()/B:.4f}, acc={acc:.2f}")

print("\nIf accuracy reaches 1.0, the full training path works.")

print("If accuracy stays low, something is broken.")

if __name__ == "__main__":

main()