Upload folder using huggingface_hub

.gitattributes

@@ -293,3 +293,13 @@ arc1-Gym2oL/step_25933 filter=lfs diff=lfs merge=lfs -text
 arc1-Gym2oL/step_259332 filter=lfs diff=lfs merge=lfs -text
 arc1-Gym2oL/step_51866 filter=lfs diff=lfs merge=lfs -text
 arc1-Gym2oL/step_77800 filter=lfs diff=lfs merge=lfs -text
+sudoku/base_mlp_sid0/step_14646 filter=lfs diff=lfs merge=lfs -text
+sudoku/base_mlp_sid0/step_19528 filter=lfs diff=lfs merge=lfs -text
+sudoku/base_mlp_sid0/step_24410 filter=lfs diff=lfs merge=lfs -text
+sudoku/base_mlp_sid0/step_29292 filter=lfs diff=lfs merge=lfs -text
+sudoku/base_mlp_sid0/step_34174 filter=lfs diff=lfs merge=lfs -text
+sudoku/base_mlp_sid0/step_39056 filter=lfs diff=lfs merge=lfs -text
+sudoku/base_mlp_sid0/step_43938 filter=lfs diff=lfs merge=lfs -text
+sudoku/base_mlp_sid0/step_4882 filter=lfs diff=lfs merge=lfs -text
+sudoku/base_mlp_sid0/step_48820 filter=lfs diff=lfs merge=lfs -text
+sudoku/base_mlp_sid0/step_9764 filter=lfs diff=lfs merge=lfs -text

sudoku/base_mlp_sid0/all_config.yaml

@@ -0,0 +1,55 @@
+arch:
+  ACT_threshold: 0
+  H_cycles: 3
+  H_layers: 0
+  L_cycles: 6
+  L_layers: 2
+  expansion: 4
+  forward_dtype: bfloat16
+  group_emb_len: 0
+  halt_exploration_prob: 0.1
+  halt_max_steps: 16
+  hidden_size: 512
+  loss:
+    alpha: 0.5
+    loss_type: stablemax_cross_entropy
+    name: losses@ACTLossHead
+  mlp_t: true
+  name: recursive_reasoning.trm@TinyRecursiveReasoningModel_ACTV1
+  no_ACT_continue: true
+  num_heads: 8
+  pos_encodings: none
+  puzzle_emb_len: 16
+  puzzle_emb_ndim: 512
+  reinit_variance: 1
+  replace_halt_threshold: 1
+  target_latent_for_reinit: L
+  universal_emb_len: 0
+beta1: 0.9
+beta2: 0.95
+checkpoint_every_eval: true
+checkpoint_path: ./checkpoint/sudoku/sudoku_base_mlp_sid0
+data_paths:
+- TRM/data/sudoku-extreme-1k-aug-1000
+data_paths_test: []
+ema: true
+ema_rate: 0.999
+epochs: 50000
+eval_interval: 5000
+eval_save_outputs: []
+evaluators: []
+freeze_weights: false
+global_batch_size: 1024
+group_emb_lr: 0.001
+group_emb_weight_decay: 0.1
+load_checkpoint: null
+lr: 0.0001
+lr_min_ratio: 1.0
+lr_warmup_steps: 2000
+min_eval_interval: 0
+project_name: pretrain_
+puzzle_emb_lr: 0.0001
+puzzle_emb_weight_decay: 1.0
+run_name: sudoku_base_mlp_sid0
+seed: 0
+weight_decay: 1.0

sudoku/base_mlp_sid0/losses.py

@@ -0,0 +1,490 @@
+from typing import Any, Tuple, Dict, Sequence, Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+import math
+
+IGNORE_LABEL_ID = -100
+
+
+def s(x, epsilon=1e-30):
+    return torch.where(
+        x<0,
+        1/(1-x+ epsilon),
+        x + 1
+    )
+
+
+def log_stablemax(x, dim=-1):
+    s_x = s(x)
+    return torch.log(s_x/torch.sum(s_x, dim=dim, keepdim=True))
+
+
+def stablemax_cross_entropy(logits, labels, ignore_index: int = -100, valid_mask=None):
+    logprobs = log_stablemax(logits.to(torch.float64), dim=-1)
+
+    if valid_mask is None:
+        valid_mask = (labels != ignore_index)
+    transformed_labels = torch.where(valid_mask, labels, 0)
+    prediction_logprobs = torch.gather(logprobs, index=transformed_labels.to(torch.long).unsqueeze(-1), dim=-1).squeeze(-1)
+
+    return -torch.where(valid_mask, prediction_logprobs, 0)
+
+def softmax_cross_entropy(logits, labels, ignore_index: int = -100):
+    # Cast logits to f32
+    # Flatten logits
+    return F.cross_entropy(logits.to(torch.float32).view(-1, logits.shape[-1]), labels.to(torch.long).view(-1), ignore_index=ignore_index, reduction="none").view(labels.shape)
+
+
+def clip_ratio_loss(cur_logits, old_logits, labels, cliprange: float = 0.1, ignore_index: int = -100, valid_mask=None):
+    logprobs = log_stablemax(cur_logits.to(torch.float64), dim=-1)
+    old_logprobs = log_stablemax(old_logits.to(torch.float64), dim=-1)
+
+    if valid_mask is None:
+        valid_mask = (labels != ignore_index)
+    transformed_labels = torch.where(valid_mask, labels, 0)
+    prediction_logprobs = torch.gather(logprobs, index=transformed_labels.to(torch.long).unsqueeze(-1), dim=-1).squeeze(-1)
+    old_prediction_logprobs = torch.gather(old_logprobs, index=transformed_labels.to(torch.long).unsqueeze(-1), dim=-1).squeeze(-1)
+
+    ratio = (prediction_logprobs - old_prediction_logprobs)
+    min_log = ratio.new_tensor(1.0 - cliprange).log()
+    max_log = ratio.new_tensor(1.0 + cliprange).log()
+    clipped_ratio = torch.clamp(ratio, min_log, max_log)
+    return -torch.where(valid_mask, torch.minimum(ratio, clipped_ratio), 0)
+
+def clip_ratio_loss2(
+    cur_logits,
+    old_logits,
+    labels,
+    advantages=None,
+    cliprange: float = 0.1,
+    ignore_index: int = -100,
+    valid_mask=None,
+):
+    if valid_mask is None:
+        valid_mask = labels != ignore_index
+
+    safe_labels = torch.where(valid_mask, labels, 0).long()
+
+    logprobs = torch.log_softmax(cur_logits.float(), dim=-1)
+
+    with torch.no_grad():
+        old_logprobs = torch.log_softmax(old_logits.float(), dim=-1)
+
+    new_logp = torch.gather(
+        logprobs,
+        dim=-1,
+        index=safe_labels.unsqueeze(-1),
+    ).squeeze(-1)
+
+    old_logp = torch.gather(
+        old_logprobs,
+        dim=-1,
+        index=safe_labels.unsqueeze(-1),
+    ).squeeze(-1)
+
+    log_ratio = new_logp - old_logp
+    ratio = torch.exp(log_ratio)
+
+    clipped_ratio = torch.clamp(
+        ratio,
+        1.0 - cliprange,
+        1.0 + cliprange,
+    )
+
+    if advantages is None:
+        advantages = torch.ones_like(ratio)
+
+    surrogate1 = ratio * advantages
+    surrogate2 = clipped_ratio * advantages
+
+    loss = -torch.minimum(surrogate1, surrogate2)
+
+    loss = torch.where(valid_mask, loss, torch.zeros_like(loss))
+
+    return loss
+
+
+def clip_ratio_loss_with_reward(
+    cur_logits,
+    old_logits,
+    labels,
+    cliprange: float = 0.1,
+    ignore_index: int = -100,
+    valid_mask=None,
+):
+    if valid_mask is None:
+        valid_mask = labels != ignore_index
+
+    with torch.no_grad():
+        pred = torch.argmax(cur_logits, dim=-1)
+        is_correct = valid_mask & (pred == labels)
+
+        reward = is_correct.sum(-1, keepdim=True).float()
+
+        # normalize by sequence length
+        denom = valid_mask.sum(-1, keepdim=True).clamp_min(1).float()
+        reward = reward / denom
+
+        reward = reward.expand_as(labels)
+        reward = torch.where(valid_mask, reward, torch.zeros_like(reward))
+
+    loss = clip_ratio_loss2(
+        cur_logits=cur_logits,
+        old_logits=old_logits,
+        labels=labels,
+        advantages=reward,
+        cliprange=cliprange,
+        ignore_index=ignore_index,
+        valid_mask=valid_mask,
+    )
+
+    return loss
+
+
+class ACTLossHead(nn.Module):
+    def __init__(self, model: nn.Module, loss_type: str, alpha: float = 0.5):
+        super().__init__()
+        self.model = model
+        self.loss_fn = globals()[loss_type]
+        
+    def initial_carry(self, *args, **kwargs):
+        return self.model.initial_carry(*args, **kwargs)  # type: ignore
+    
+    @staticmethod
+    def comput_q_halt_loss(self, outputs, is_correct, seq_is_correct, loss_counts):
+        q_halt_loss = F.binary_cross_entropy_with_logits(outputs["q_halt_logits"], seq_is_correct.to(outputs["q_halt_logits"].dtype), reduction="sum")
+        return q_halt_loss
+
+    def compute_lm_loss(self, outputs, carry, labels, mask):
+        return self.loss_fn(outputs["logits"], labels, ignore_index=IGNORE_LABEL_ID, valid_mask=mask)
+    
+    def update_carry(self, carry, outputs):
+        return carry 
+    
+    def forward(
+        self,
+        return_keys: Sequence[str],
+        # Model args
+        **model_kwargs,
+    ) -> Tuple[Any, torch.Tensor, Dict[str, torch.Tensor], Optional[Dict[str, torch.Tensor]], torch.Tensor]:
+        # Model logits
+        # B x SeqLen x D
+        new_carry, outputs = self.model(**model_kwargs)
+        labels = new_carry.current_data["labels"]
+
+        with torch.no_grad():
+            # Preds
+            outputs["preds"] = torch.argmax(outputs["logits"], dim=-1)
+
+            # Correctness
+            mask = (labels != IGNORE_LABEL_ID)
+            loss_counts = mask.sum(-1)
+            loss_divisor = loss_counts.clamp_min(1).unsqueeze(-1)  # Avoid NaNs in division
+
+            is_correct = mask & (torch.argmax(outputs["logits"], dim=-1) == labels)
+            seq_is_correct = is_correct.sum(-1) == loss_counts
+            
+            # Metrics (halted)
+            valid_metrics = new_carry.halted & (loss_counts > 0)
+            metrics = {
+                "count": valid_metrics.sum(),
+                
+                "accuracy":       torch.where(valid_metrics, (is_correct.to(torch.float32) / loss_divisor).sum(-1), 0).sum(),
+                "exact_accuracy": (valid_metrics & seq_is_correct).sum(),
+
+                "q_halt_accuracy": (valid_metrics & ((outputs["q_halt_logits"] >= 0) == seq_is_correct)).sum(),
+                "steps":          torch.where(valid_metrics, new_carry.steps, 0).sum(),
+            }
+
+        # Losses
+        lm_loss = (self.compute_lm_loss(outputs, new_carry, labels, mask) / loss_divisor).sum()
+        q_halt_loss = self.comput_q_halt_loss(self, outputs, is_correct, seq_is_correct, loss_counts)
+        new_carry = self.update_carry(new_carry, outputs)
+        metrics.update({
+            "lm_loss": lm_loss.detach(),
+            "q_halt_loss": q_halt_loss.detach(),
+        })
+        # Q continue (bootstrapping target loss); Alexia: This fits Q-learning, but seems totally unecessary
+        q_continue_loss = 0
+        if "target_q_continue" in outputs:
+            q_continue_loss = F.binary_cross_entropy_with_logits(outputs["q_continue_logits"], outputs["target_q_continue"], reduction="sum")
+
+            metrics["q_continue_loss"] = q_continue_loss.detach()
+        # Filter outputs for return
+        detached_outputs = {k: outputs[k].detach() for k in return_keys if k in outputs}
+
+        return new_carry, lm_loss + 0.5 * (q_halt_loss + q_continue_loss), metrics, detached_outputs, new_carry.halted.all()
+
+class ACTLossHead_binorminal(ACTLossHead):
+    def __init__(self, model: nn.Module, loss_type: str, alpha: float = 0.5):
+        super().__init__(model, loss_type)
+        
+    @staticmethod
+    def comput_q_halt_loss(self, outputs, is_correct, seq_is_correct, loss_counts):
+        t = (is_correct.sum(-1) / loss_counts.clamp_min(1)).to(outputs["q_halt_logits"].dtype)  # [B] in [0,1]
+
+        # Force t=1 for sequences with no valid tokens
+        t = torch.where(loss_counts == 0, torch.ones_like(t), t)
+        q_halt_loss = F.binary_cross_entropy_with_logits(
+            outputs["q_halt_logits"], t, reduction="sum"
+        )        
+        return q_halt_loss
+
+class ACTLossHead_combine(ACTLossHead):
+    def __init__(self, model: nn.Module, loss_type: str, alpha: float = 0.5):
+        super().__init__(model, loss_type)
+        self.alpha = alpha
+
+    @staticmethod
+    def comput_q_halt_loss(self, outputs, is_correct, seq_is_correct, loss_counts):
+        q_halt_bce_loss = ACTLossHead.comput_q_halt_loss(self, outputs, is_correct, seq_is_correct, loss_counts)
+        q_halt_binorm_loss = ACTLossHead_binorminal.comput_q_halt_loss(self, outputs, is_correct, seq_is_correct, loss_counts)
+        q_halt_loss = self.alpha * q_halt_bce_loss + (1 - self.alpha) * q_halt_binorm_loss
+        return q_halt_loss
+
+class ACTLossHead_binormonly(ACTLossHead_binorminal):
+    def __init__(self, model: nn.Module, loss_type: str):
+        super().__init__(model, loss_type)
+
+    def forward(
+        self,
+        return_keys: Sequence[str],
+        # Model args
+        **model_kwargs,
+    ) -> Tuple[Any, torch.Tensor, Dict[str, torch.Tensor], Optional[Dict[str, torch.Tensor]], torch.Tensor]:
+        # Model logits
+        # B x SeqLen x D
+        new_carry, outputs = self.model(**model_kwargs)
+        labels = new_carry.current_data["labels"]
+
+        with torch.no_grad():
+            # Preds
+            outputs["preds"] = torch.argmax(outputs["logits"], dim=-1)
+
+            # Correctness
+            mask = (labels != IGNORE_LABEL_ID)
+            loss_counts = mask.sum(-1)
+            loss_divisor = loss_counts.clamp_min(1).unsqueeze(-1)  # Avoid NaNs in division
+
+            is_correct = mask & (torch.argmax(outputs["logits"], dim=-1) == labels)
+            seq_is_correct = is_correct.sum(-1) == loss_counts
+            
+            # Metrics (halted)
+            valid_metrics = new_carry.halted & (loss_counts > 0)
+            metrics = {
+                "count": valid_metrics.sum(),
+                
+                "accuracy":       torch.where(valid_metrics, (is_correct.to(torch.float32) / loss_divisor).sum(-1), 0).sum(),
+                "exact_accuracy": (valid_metrics & seq_is_correct).sum(),
+
+                "q_halt_accuracy": (valid_metrics & ((outputs["q_halt_logits"] >= 0) == seq_is_correct)).sum(),
+                "steps":          torch.where(valid_metrics, new_carry.steps, 0).sum(),
+            }
+
+        # Losses
+        lm_loss = (self.compute_lm_loss(outputs, new_carry, labels, mask) / loss_divisor).sum()
+        q_halt_loss = self.comput_q_halt_loss(self, outputs, is_correct, seq_is_correct, loss_counts)
+        metrics.update({
+            "lm_loss": lm_loss.detach(),
+            "q_halt_loss": q_halt_loss.detach(),
+        })
+        # Q continue (bootstrapping target loss); Alexia: This fits Q-learning, but seems totally unecessary
+        q_continue_loss = 0
+        if "target_q_continue" in outputs:
+            q_continue_loss = F.binary_cross_entropy_with_logits(outputs["q_continue_logits"], outputs["target_q_continue"], reduction="sum")
+
+            metrics["q_continue_loss"] = q_continue_loss.detach()
+        # Filter outputs for return
+        detached_outputs = {k: outputs[k].detach() for k in return_keys if k in outputs}
+
+        return new_carry,  (q_halt_loss + q_continue_loss), metrics, detached_outputs, new_carry.halted.all()
+
+class TRMPPOLossHead(nn.Module):
+    def __init__(
+        self,
+        model: nn.Module,
+        cliprange: float = 0.1,
+        ent_coef: float = 0.01,
+        vf_coef: float = 0.5,
+        gamma: float = 0.99,
+        lam: float = 0.95,
+        ppo_noptepochs: int = 4,
+        reward_type: str = "sparse",
+    ) -> None:
+        super().__init__()
+        self.model = model
+        self.cliprange = cliprange
+        self.ent_coef = ent_coef
+        self.vf_coef = vf_coef
+        self.gamma = gamma
+        self.lam = lam
+        self.ppo_noptepochs = ppo_noptepochs
+        self.reward_type = reward_type
+
+    def initial_carry(self, *args, **kwargs):
+        return self.model.initial_carry(*args, **kwargs)  # type: ignore
+
+    def _compute_ppo_signal(
+        self,
+        labels: torch.Tensor,
+        trajectory_logits: torch.Tensor,
+        trajectory_values: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        trajectory_values = trajectory_values.squeeze(-1)
+        mask = labels != IGNORE_LABEL_ID
+        mask_exp = mask.unsqueeze(1).expand(-1, trajectory_logits.shape[1], -1)
+        valid_row_mask = mask.sum(-1) > 0  # [B] - True if row has valid tokens
+        valid_adv_mask = valid_row_mask.unsqueeze(-1).expand(-1, trajectory_logits.shape[1])  # [B, T]
+
+        preds = torch.argmax(trajectory_logits, dim=-1)
+        is_correct = mask_exp & (preds == labels.unsqueeze(1))
+
+        loss_counts = mask.sum(-1).unsqueeze(1)
+        seq_is_correct = is_correct.sum(-1) == loss_counts
+        seq_is_correct_f = seq_is_correct.to(torch.float32)
+
+        value_pred = trajectory_values
+
+        # Assign step-specific rewards and filter out padded rows
+        rewards = torch.zeros_like(value_pred) # [B, T]
+        if self.reward_type == "sparse":
+            rewards[:, -1] = seq_is_correct_f[:, -1] * valid_row_mask
+        elif self.reward_type == "dense":
+            rewards = seq_is_correct_f * valid_row_mask.unsqueeze(-1)  # [B, T]
+        else:
+            raise ValueError(f"Unsupported reward type: {self.reward_type}")
+
+        next_values = torch.zeros_like(value_pred)
+        next_values[:, :-1] = value_pred[:, 1:]
+
+        deltas = rewards + self.gamma * next_values - value_pred
+
+        advs = torch.zeros_like(deltas)
+        last_gae = torch.zeros_like(deltas[:, 0])
+        for t in range(deltas.shape[1] - 1, -1, -1):
+            last_gae = deltas[:, t] + self.gamma * self.lam * last_gae
+            advs[:, t] = last_gae
+
+        # Returns based on end-of-trajectory reward
+        # This is likely to be appropriate with sparse rewards 
+        returns = torch.zeros_like(value_pred)
+        final_reward = rewards[:, -1]  # Shape: [B]
+
+        for t in range(deltas.shape[1]):
+            distance_to_end = deltas.shape[1] - 1 - t
+            returns[:, t] = final_reward * (self.gamma ** distance_to_end)
+
+        return mask, mask_exp, valid_adv_mask, rewards, advs, returns, value_pred, seq_is_correct, is_correct, loss_counts
+
+    def compute_ppo_baseline_returns(
+        self,
+        labels: torch.Tensor,
+        trajectory_logits: torch.Tensor,
+        trajectory_values: torch.Tensor,
+    ) -> torch.Tensor:
+        _, _, _, _, advs, _, value_pred, _, _, _ = self._compute_ppo_signal(
+            labels, trajectory_logits, trajectory_values
+        )
+        return advs + value_pred
+
+    def forward(
+        self,
+        return_keys: Sequence[str],
+        **model_kwargs,
+    ) -> Tuple[Any, torch.Tensor, Dict[str, torch.Tensor], Optional[Dict[str, torch.Tensor]], torch.Tensor]:
+        new_carry, outputs = self.model(**model_kwargs)
+        labels = new_carry.current_data["labels"]
+
+        trajectory_logits = outputs["trajectory_logits"]
+        trajectory_values = outputs["trajectory_values"]
+
+        # Replace ignore labels so gather stays in-bounds; mask applied later.
+        safe_labels = torch.where(labels == IGNORE_LABEL_ID, torch.zeros_like(labels), labels)
+        labels_exp = safe_labels.unsqueeze(1).expand(-1, trajectory_logits.shape[1], -1).unsqueeze(-1)
+        logprobs = F.log_softmax(trajectory_logits.to(torch.float32), dim=-1)
+        action_logprobs = torch.gather(logprobs, dim=-1, index=labels_exp.to(torch.long)).squeeze(-1)
+        mask, mask_exp, valid_adv_mask, rewards, advs, returns, value_pred, seq_is_correct, is_correct, loss_counts = self._compute_ppo_signal(
+            labels, trajectory_logits, trajectory_values
+        )
+
+        action_logprobs = torch.where(mask_exp, action_logprobs, torch.zeros_like(action_logprobs))
+
+        batch = model_kwargs.get("batch")
+        returns = None
+        if isinstance(batch, dict) and batch.get("ppo_target_returns") is not None:
+            returns = batch["ppo_target_returns"]
+
+        assert returns is not None, "PPO target returns must be provided in the batch for the PPO loss head."
+
+        valid_advs = advs[valid_adv_mask]
+        if valid_advs.numel() > 0:
+            # Calculate mean and variance strictly over real data points
+            adv_mean = valid_advs.mean()
+            adv_std = valid_advs.std().clamp_min(1e-8)
+            
+            # Normalize the entire tensor using the clean statistics
+            advs = (advs - adv_mean) / adv_std
+        else:
+            # Fallback to prevent NaNs if an empty batch somehow passes through
+            advs = torch.zeros_like(advs)
+
+        # 3. Explicitly zero-out advantages on padded rows to prevent gradient updates
+        advs = torch.where(valid_adv_mask, advs, torch.zeros_like(advs))
+
+        old_action_logprobs = None
+        old_action_logprobs = batch.get("ppo_old_action_logprobs")
+        if old_action_logprobs is None:
+            old_action_logprobs = action_logprobs.detach()
+        old_action_logprobs = torch.where(mask_exp, old_action_logprobs, torch.zeros_like(old_action_logprobs))
+
+        token_ratio = torch.exp(action_logprobs - old_action_logprobs)
+        clipped_token_ratio = torch.clamp(token_ratio, 1.0 - self.cliprange, 1.0 + self.cliprange)
+        advs_expanded = advs.unsqueeze(-1)
+        surr1 = token_ratio * advs_expanded
+        surr2 = clipped_token_ratio * advs_expanded
+        token_pg_loss = -torch.minimum(surr1, surr2)
+        token_pg_loss = torch.where(mask_exp, token_pg_loss, torch.zeros_like(token_pg_loss))
+        pg_loss = token_pg_loss.sum() / mask_exp.sum().clamp_min(1)
+
+        value_loss = 0.5 * torch.mean((value_pred - returns) ** 2)
+
+        probs = logprobs.exp()
+        token_entropy = -(probs * logprobs).sum(-1)
+        entropy = (token_entropy * mask_exp).sum() / mask_exp.sum().clamp_min(1)
+
+        loss = pg_loss + self.vf_coef * value_loss - self.ent_coef * entropy
+
+
+        valid_metrics = new_carry.halted & (loss_counts > 0)
+        loss_divisor = loss_counts.clamp_min(1).unsqueeze(-1)  # Avoid NaNs in division
+        metrics = {
+            "pg_loss": pg_loss.detach(),
+            "value_loss": value_loss.detach(),
+            "entropy": entropy.detach(),
+            "reward_mean": rewards.mean().detach(),
+            "count": valid_metrics.sum(),
+            "accuracy":       torch.where(valid_metrics, (is_correct[:, -1].to(torch.float32) / loss_divisor).sum(-1), 0).sum(),
+            "exact_accuracy": (valid_metrics & seq_is_correct[:, -1]).sum(),
+            "steps":          torch.where(valid_metrics, new_carry.steps, 0).sum(),
+        }
+
+
+        outputs["ppo_old_logprobs"] = action_logprobs.sum(-1).detach()
+        detached_outputs = {k: outputs[k].detach() for k in return_keys if k in outputs}
+
+        return new_carry, loss, metrics, detached_outputs, new_carry.halted.all()
+
+class ACTClipRatioLossHead(ACTLossHead):
+    def __init__(self, model: nn.Module, loss_type: str = "clip_ratio_loss", cliprange: float = 0.1):
+        super().__init__(model, loss_type)
+        self.cliprange = cliprange
+
+    def compute_lm_loss(self, outputs, carry, labels, mask):
+        lm_loss = self.loss_fn(outputs["logits"], carry.oldlogits, labels, cliprange=self.cliprange, ignore_index=IGNORE_LABEL_ID, valid_mask=mask)
+        return lm_loss
+
+    def update_carry(self, carry, outputs):
+        carry.oldlogits = outputs["logits"].detach()
+        return carry

sudoku/base_mlp_sid0/step_14646

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

sudoku/base_mlp_sid0/step_19528

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

sudoku/base_mlp_sid0/step_24410

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

sudoku/base_mlp_sid0/step_29292

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

sudoku/base_mlp_sid0/step_34174

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

sudoku/base_mlp_sid0/step_39056

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

sudoku/base_mlp_sid0/step_43938

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

sudoku/base_mlp_sid0/step_4882

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

sudoku/base_mlp_sid0/step_48820

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

sudoku/base_mlp_sid0/step_9764

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

sudoku/base_mlp_sid0/trm.py

@@ -0,0 +1,645 @@
+from typing import Tuple, List, Dict, Optional
+from dataclasses import dataclass
+import math
+import torch
+import copy
+import torch.nn.functional as F
+from torch import nn
+from pydantic import BaseModel
+import random
+from models.common import trunc_normal_init_
+from models.layers import rms_norm, LinearSwish, SwiGLU, Attention, RotaryEmbedding, CosSin, CastedEmbedding, CastedLinear
+from models.sparse_embedding import CastedSparseEmbedding
+
+IGNORE_LABEL_ID = -100
+
+@dataclass
+class TinyRecursiveReasoningModel_ACTV1InnerCarry:
+    z_H: torch.Tensor
+    z_L: torch.Tensor
+
+@dataclass
+class TinyRecursiveReasoningModel_ACTV1Carry:
+    inner_carry: TinyRecursiveReasoningModel_ACTV1InnerCarry
+    
+    steps: torch.Tensor
+    halted: torch.Tensor
+    
+    current_data: Dict[str, torch.Tensor]
+
+@dataclass
+class TinyRecursiveReasoningModel_ACTVClipCarry:
+    inner_carry: TinyRecursiveReasoningModel_ACTV1InnerCarry
+    
+    oldlogits: torch.Tensor
+
+    steps: torch.Tensor
+    halted: torch.Tensor
+    
+    current_data: Dict[str, torch.Tensor]
+
+
+class TinyRecursiveReasoningModel_ACTV1Config(BaseModel):
+    batch_size: int
+    seq_len: int
+    puzzle_emb_ndim: int = 0
+    universal_emb_len: int = 0
+    group_emb_len: int = 0
+    num_puzzle_identifiers: int
+    num_groups: int
+    vocab_size: int
+
+    H_cycles: int
+    L_cycles: int
+
+    H_layers: int # ignored
+    L_layers: int
+
+    # Transformer config
+    hidden_size: int
+    expansion: float
+    num_heads: int
+    pos_encodings: str
+
+    rms_norm_eps: float = 1e-5
+    rope_theta: float = 10000.0
+    
+    # Halting Q-learning config
+    halt_max_steps: int
+    halt_exploration_prob: float
+    puzzle_emb_len: int = 16 # if non-zero, its specified to this value
+    forward_dtype: str = "bfloat16"
+
+    # Alexia: added
+    mlp_t: bool = False # use mlp on L instead of transformer
+    no_ACT_continue: bool =  True # No continue ACT loss, only use the sigmoid of the halt which makes much more sense
+
+    # chan edit
+    eval_with_ACT: bool = False # During evaluation, whether to use ACT halting or just run max steps. Using ACT halting will cause more variance in evaluation results, but is a more realistic evaluation of the model performance when deployed.
+    ACT_threshold: float = 0 # Threshold for ACT halting during evaluation, only used if eval_with_ACT is True
+
+class TinyRecursiveReasoningModel_ACTGymConfig(BaseModel):
+    replace_halt_threshold: int = 1
+    target_latent_for_reinit: str = "L"
+    reinit_variance: float = 1.0 # Variance for reinitializing the latent when halted. This is a hyperparameter that can be tuned. Setting it to 0 will reinitialize to the same value, which may cause the model to learn to rely on this specific value and not learn effectively. Setting it to a very large value will make the reinitialized latent too random, which may also harm learning. A moderate value can help the model learn to recover from halting and continue learning effectively.
+
+
+class TinyRecursiveReasoningModel_ACTV1Block(nn.Module):
+    def __init__(self, config: TinyRecursiveReasoningModel_ACTV1Config, puzzle_emb_len: int) -> None:
+        super().__init__()
+
+        self.config = config
+        if self.config.mlp_t:
+            self.puzzle_emb_len = puzzle_emb_len
+            self.mlp_t = SwiGLU(
+                hidden_size=self.config.seq_len + self.puzzle_emb_len, # L
+                expansion=config.expansion,
+            )
+        else:
+            self.self_attn = Attention(
+                hidden_size=config.hidden_size,
+                head_dim=config.hidden_size // config.num_heads,
+                num_heads=config.num_heads,
+                num_key_value_heads=config.num_heads,
+                causal=False
+            )
+        self.mlp = SwiGLU(
+            hidden_size=config.hidden_size,
+            expansion=config.expansion,
+        )
+        self.norm_eps = config.rms_norm_eps
+
+    def forward(self, cos_sin: CosSin, hidden_states: torch.Tensor) -> torch.Tensor:
+        # B, L, D = hidden_states.shape
+        # Post Norm
+        if self.config.mlp_t:
+            hidden_states = hidden_states.transpose(1,2)
+            out = self.mlp_t(hidden_states)
+            hidden_states = rms_norm(hidden_states + out, variance_epsilon=self.norm_eps)
+            hidden_states = hidden_states.transpose(1,2)
+        else:
+            # Self Attention
+            hidden_states = rms_norm(hidden_states + self.self_attn(cos_sin=cos_sin, hidden_states=hidden_states), variance_epsilon=self.norm_eps)
+        # Fully Connected
+        out = self.mlp(hidden_states)
+        hidden_states = rms_norm(hidden_states + out, variance_epsilon=self.norm_eps)
+        return hidden_states
+
+class TinyRecursiveReasoningModel_ACTV1ReasoningModule(nn.Module):
+    def __init__(self, layers: List[TinyRecursiveReasoningModel_ACTV1Block]):
+        super().__init__()
+        self.layers = torch.nn.ModuleList(layers)
+
+    def forward(self, hidden_states: torch.Tensor, input_injection: torch.Tensor, **kwargs) -> torch.Tensor:
+        hidden_states = hidden_states + input_injection
+        for layer in self.layers:
+            hidden_states = layer(hidden_states=hidden_states, **kwargs)
+        return hidden_states
+
+class TinyRecursiveReasoningModel_ACTV1_Inner(nn.Module):
+    def __init__(self, config: TinyRecursiveReasoningModel_ACTV1Config) -> None:
+        super().__init__()
+        self.config = config
+        self.forward_dtype = getattr(torch, self.config.forward_dtype)
+
+        # I/O
+
+        self.embed_scale = math.sqrt(self.config.hidden_size)
+        embed_init_std = 1.0 / self.embed_scale
+
+        self.embed_tokens = CastedEmbedding(self.config.vocab_size, self.config.hidden_size, init_std=embed_init_std, cast_to=self.forward_dtype)
+        self.lm_head      = CastedLinear(self.config.hidden_size, self.config.vocab_size, bias=False)
+        self.q_head       = CastedLinear(self.config.hidden_size, 2, bias=True)
+
+        if self.config.puzzle_emb_ndim > 0:
+            # Zero init puzzle embeddings
+            self.puzzle_emb = CastedSparseEmbedding(self.config.num_puzzle_identifiers, self.config.puzzle_emb_ndim,
+                                                    batch_size=self.config.batch_size, init_std=0, cast_to=self.forward_dtype)
+        
+        self.group_embedding_dim = self.config.group_emb_len * self.config.hidden_size
+        if self.group_embedding_dim > 0:
+            # Zero init group embeddings
+            self.group_emb = CastedSparseEmbedding(self.config.num_groups, self.group_embedding_dim,
+                                                    batch_size=self.config.batch_size, init_std=0, cast_to=self.forward_dtype)
+
+        self.universal_emb_ndim = self.config.universal_emb_len * self.config.hidden_size
+        if self.universal_emb_ndim > 0: 
+            # Universal embedding, shared between tokens and puzzles, with non-zero init
+            self.universal_emb = CastedEmbedding(1, self.universal_emb_ndim, init_std=embed_init_std, cast_to=self.forward_dtype)
+        
+        print("Puzzle embedding ndim:", self.config.puzzle_emb_ndim)
+        print("Group embedding ndim:", self.group_embedding_dim)
+        print("Universal embedding ndim:", self.universal_emb_ndim)
+        print("Puzzle embed length:", self.config.puzzle_emb_len)
+        total_puzzle_emb_dim = self.config.puzzle_emb_ndim + self.group_embedding_dim + self.universal_emb_ndim
+        self.puzzle_emb_len = max(-(total_puzzle_emb_dim // -self.config.hidden_size),  self.config.puzzle_emb_len)  # ceil div
+        print("Total puzzle embedding dim:", total_puzzle_emb_dim)
+        print("Final puzzle embedding length (in tokens):", self.puzzle_emb_len)
+
+        # LM Blocks
+        if self.config.pos_encodings == "rope":
+            self.rotary_emb = RotaryEmbedding(dim=self.config.hidden_size // self.config.num_heads,
+                                              max_position_embeddings=self.config.seq_len + self.puzzle_emb_len,
+                                              base=self.config.rope_theta)
+        elif self.config.pos_encodings == "learned":
+            self.embed_pos = CastedEmbedding(self.config.seq_len + self.puzzle_emb_len, self.config.hidden_size, init_std=embed_init_std, cast_to=self.forward_dtype)
+        else:
+            pass
+
+        # Reasoning Layers
+        self.L_level = TinyRecursiveReasoningModel_ACTV1ReasoningModule(layers=[TinyRecursiveReasoningModel_ACTV1Block(self.config, self.puzzle_emb_len) for _i in range(self.config.L_layers)])
+
+        # Initial states
+        self.H_init = self.init_latent_buffer(self.config.hidden_size, self.forward_dtype)
+        self.L_init = self.init_latent_buffer(self.config.hidden_size, self.forward_dtype)
+
+        # Q head special init
+        # Init Q to (almost) zero for faster learning during bootstrapping
+        with torch.no_grad():
+            self.q_head.weight.zero_()
+            self.q_head.bias.fill_(-5)  # type: ignore
+
+    @staticmethod
+    def init_latent_buffer(
+        hidden_size: int,
+        forward_dtype: torch.dtype,
+        persistent: bool = True,
+        device: Optional[torch.device] = None,
+        std: float = 1.0
+    ) -> torch.Tensor:
+        return nn.Buffer(
+            trunc_normal_init_(torch.empty(hidden_size, dtype=forward_dtype, device=device), std=std),
+            persistent=persistent,
+        )
+        
+    
+    @staticmethod
+    def init_latent_buffer_around_state(
+        state: torch.Tensor,
+        forward_dtype: torch.dtype,
+        persistent: bool = True,    
+        device: Optional[torch.device] = None,
+        std: float = 1.0
+    ) -> torch.Tensor:
+        return nn.Buffer(
+            state + trunc_normal_init_(torch.empty_like(state, dtype=forward_dtype, device=device), std=std),
+            persistent=persistent,
+        )
+    
+
+
+    def _input_embeddings(self, input: torch.Tensor, puzzle_identifiers: torch.Tensor, group_indices: torch.Tensor):
+        # Token embedding
+        embedding = self.embed_tokens(input.to(torch.int32))
+
+        # Puzzle embeddings
+        if self.config.puzzle_emb_ndim > 0:
+            puzzle_embedding = self.puzzle_emb(puzzle_identifiers)
+
+            if self.config.group_emb_len > 0:
+                group_embedding = self.group_emb(group_indices)
+
+                # concat along feature dim -> (B, Dp + Dg)
+                puzzle_embedding = torch.cat([puzzle_embedding, group_embedding], dim=-1)
+
+            if self.config.universal_emb_len > 0:
+                # universal_embedding: (B, Du)
+                universal_embedding = self.universal_emb(torch.zeros(puzzle_embedding.size(0), device=puzzle_embedding.device, dtype=torch.long))
+
+                # concat along feature dim -> (B, Dp + (Dg) + Du)
+                puzzle_embedding = torch.cat([puzzle_embedding, universal_embedding], dim=-1)
+                
+            pad_count = self.puzzle_emb_len * self.config.hidden_size - puzzle_embedding.shape[-1]
+            if pad_count > 0:
+                puzzle_embedding = F.pad(puzzle_embedding, (0, pad_count))
+
+            embedding = torch.cat((puzzle_embedding.view(-1, self.puzzle_emb_len, self.config.hidden_size), embedding), dim=-2)
+           
+        # Position embeddings
+        if self.config.pos_encodings == "learned":
+            # scale by 1/sqrt(2) to maintain forward variance
+            embedding = 0.707106781 * (embedding + self.embed_pos.embedding_weight.to(self.forward_dtype))
+
+        # Scale
+        return self.embed_scale * embedding
+
+    def empty_carry(self, batch_size: int):
+        return TinyRecursiveReasoningModel_ACTV1InnerCarry(
+            z_H=torch.empty(batch_size, self.config.seq_len + self.puzzle_emb_len, self.config.hidden_size, dtype=self.forward_dtype),
+            z_L=torch.empty(batch_size, self.config.seq_len + self.puzzle_emb_len, self.config.hidden_size, dtype=self.forward_dtype),
+        )
+        
+    def reset_carry(self, reset_flag: torch.Tensor, carry: TinyRecursiveReasoningModel_ACTV1InnerCarry):
+        return TinyRecursiveReasoningModel_ACTV1InnerCarry(
+            z_H=torch.where(reset_flag.view(-1, 1, 1), self.H_init, carry.z_H),
+            z_L=torch.where(reset_flag.view(-1, 1, 1), self.L_init, carry.z_L),
+        )
+
+    def re_inittialize_carry(
+            self, 
+            reset_flag: torch.Tensor, 
+            carry: TinyRecursiveReasoningModel_ACTV1InnerCarry, 
+            target_latent: Optional[str] = None,
+            init_std: float = 1.0
+        ):
+        target_device = carry.z_L.device
+        if target_latent == "LH":
+            return TinyRecursiveReasoningModel_ACTV1InnerCarry(
+                z_H=torch.where(
+                    reset_flag.view(-1, 1, 1),
+                    self.init_latent_buffer(self.config.hidden_size, self.forward_dtype, device=target_device, std=init_std),
+                    carry.z_H,
+                ),
+                z_L=torch.where(
+                    reset_flag.view(-1, 1, 1),
+                    self.init_latent_buffer(self.config.hidden_size, self.forward_dtype, device=target_device, std=init_std),
+                    carry.z_L,
+                ),
+            )
+        elif target_latent == "H":
+            return TinyRecursiveReasoningModel_ACTV1InnerCarry(
+                z_H=torch.where(
+                    reset_flag.view(-1, 1, 1),
+                    self.init_latent_buffer(self.config.hidden_size, self.forward_dtype, device=target_device, std=init_std),
+                    carry.z_H,
+                ),
+                z_L=carry.z_L
+            )
+        elif target_latent == "L":
+            return TinyRecursiveReasoningModel_ACTV1InnerCarry(
+                z_H=carry.z_H,
+                z_L=torch.where(
+                    reset_flag.view(-1, 1, 1),
+                    self.init_latent_buffer(self.config.hidden_size, self.forward_dtype, device=target_device, std=init_std),
+                    carry.z_L,
+                )
+            )
+        else:
+            raise ValueError(f"Invalid target_latent value: {target_latent}")
+
+
+    def forward(self, carry: TinyRecursiveReasoningModel_ACTV1InnerCarry, batch: Dict[str, torch.Tensor]) -> Tuple[TinyRecursiveReasoningModel_ACTV1InnerCarry, Dict[str, torch.Tensor]]:
+        seq_info = dict(
+            cos_sin=self.rotary_emb() if hasattr(self, "rotary_emb") else None,
+        )
+
+        # Input encoding
+        input_embeddings = self._input_embeddings(batch["inputs"], batch["puzzle_identifiers"], batch["group_indices"])
+
+        # Forward iterations
+        it = 0
+        z_H, z_L = carry.z_H, carry.z_L
+        # H_cycles-1 without grad
+        with torch.no_grad():
+            for _H_step in range(self.config.H_cycles-1):
+                for _L_step in range(self.config.L_cycles):
+                    z_L = self.L_level(z_L, z_H + input_embeddings, **seq_info)
+                z_H = self.L_level(z_H, z_L, **seq_info)
+        # 1 with grad
+        for _L_step in range(self.config.L_cycles):
+            z_L = self.L_level(z_L, z_H + input_embeddings, **seq_info)
+        z_H = self.L_level(z_H, z_L, **seq_info)
+
+        # LM Outputs
+        new_carry = TinyRecursiveReasoningModel_ACTV1InnerCarry(z_H=z_H.detach(), z_L=z_L.detach())  # New carry no grad
+        logits = self.lm_head(z_H)[:, self.puzzle_emb_len:]
+        q_logits = self.q_head(z_H[:, 0]).to(torch.float32) # Q-head; uses the first puzzle_emb position
+        outputs = {
+            "logits": logits,
+            "q_halt_logits": q_logits[..., 0],
+            "q_continue_logits": q_logits[..., 1] if not self.config.no_ACT_continue else None
+        }
+        
+        return new_carry, outputs  
+
+class TinyRecursiveReasoningModel_PPOInner(TinyRecursiveReasoningModel_ACTV1_Inner):
+    def __init__(self, config: TinyRecursiveReasoningModel_ACTV1Config) -> None:
+        super().__init__(config)
+        self.v_head = CastedLinear(self.config.hidden_size, 1, bias=True) # Unbounded PPO value function regression
+
+        with torch.no_grad():
+            self.v_head.weight.normal_(std=0.02)
+            self.v_head.bias.zero_()     # Clean zero-start for reward expectation
+
+    def forward(
+        self,
+        carry: TinyRecursiveReasoningModel_ACTV1InnerCarry,
+        batch: Dict[str, torch.Tensor],
+    ) -> Tuple[TinyRecursiveReasoningModel_ACTV1InnerCarry, Dict[str, torch.Tensor]]:
+        seq_info = dict(
+            cos_sin=self.rotary_emb() if hasattr(self, "rotary_emb") else None,
+        )
+
+        input_embeddings = self._input_embeddings(batch["inputs"], batch["puzzle_identifiers"], batch["group_indices"])
+
+        z_H, z_L = carry.z_H, carry.z_L
+        trajectory_logits = []
+        trajectory_values = []
+
+        with torch.no_grad():
+            for _H_step in range(self.config.H_cycles - 1):
+                for _L_step in range(self.config.L_cycles):
+                    z_L = self.L_level(z_L, z_H + input_embeddings, **seq_info)
+                z_H = self.L_level(z_H, z_L, **seq_info)
+
+                logits = self.lm_head(z_H)[:, self.puzzle_emb_len:]
+                values = self.v_head(z_H[:, 0]).to(torch.float32)
+                trajectory_logits.append(logits.detach())
+                trajectory_values.append(values.detach())
+
+        for _L_step in range(self.config.L_cycles):
+            z_L = self.L_level(z_L, z_H + input_embeddings, **seq_info)
+        z_H = self.L_level(z_H, z_L, **seq_info)
+
+        logits = self.lm_head(z_H)[:, self.puzzle_emb_len:]
+        # q_logits = self.q_head(z_H[:, 0]).to(torch.float32)
+        values = self.v_head(z_H[:, 0]).to(torch.float32)
+        trajectory_logits.append(logits)
+        trajectory_values.append(values)
+
+        new_carry = TinyRecursiveReasoningModel_ACTV1InnerCarry(z_H=z_H.detach(), z_L=z_L.detach())
+        outputs = {
+            "logits": logits,
+            "trajectory_logits": torch.stack(trajectory_logits, dim=1),
+            "trajectory_values": torch.stack(trajectory_values, dim=1)
+        }
+
+        return new_carry, outputs
+
+
+class TinyRecursiveReasoningModel_ACTV1(nn.Module):
+    """ACT wrapper."""
+
+    def __init__(self, config_dict: dict):
+        super().__init__()
+        self.config = TinyRecursiveReasoningModel_ACTV1Config(**config_dict)
+        self.inner = TinyRecursiveReasoningModel_ACTV1_Inner(self.config)
+
+    @property
+    def puzzle_emb(self):
+        return self.inner.puzzle_emb
+
+    @property
+    def group_emb(self):
+        return self.inner.group_emb
+
+    def initial_carry(self, batch: Dict[str, torch.Tensor]):
+        batch_size = batch["inputs"].shape[0]
+
+        return TinyRecursiveReasoningModel_ACTV1Carry(
+            inner_carry=self.inner.empty_carry(batch_size),  # Empty is expected, it will be reseted in first pass as all sequences are halted.
+            
+            steps=torch.zeros((batch_size, ), dtype=torch.int32),
+            halted=torch.ones((batch_size, ), dtype=torch.bool),  # Default to halted
+            
+            current_data={k: torch.empty_like(v) for k, v in batch.items()}
+        )
+
+    def step_inner_carry(self, carry: TinyRecursiveReasoningModel_ACTV1Carry, batch: Dict[str, torch.Tensor]):
+        new_inner_carry = self.inner.reset_carry(carry.halted, carry.inner_carry)
+        return new_inner_carry
+
+    def step_current_data(self, carry: TinyRecursiveReasoningModel_ACTV1Carry, batch: Dict[str, torch.Tensor]):
+        new_current_data = {k: torch.where(carry.halted.view((-1, ) + (1, ) * (batch[k].ndim - 1)), batch[k], v) for k, v in carry.current_data.items()}
+        return new_current_data
+    
+    def step_carry(self, carry: TinyRecursiveReasoningModel_ACTV1Carry, batch: Dict[str, torch.Tensor]):
+        # Step inner carry
+        new_inner_carry = self.step_inner_carry(carry, batch)
+        
+        new_steps = torch.where(carry.halted, 0, carry.steps)
+
+        # Update data, carry (removing halted sequences)        
+        new_current_data = self.step_current_data(carry, batch)
+
+        return new_inner_carry, new_steps, new_current_data
+
+
+    def forward(self, carry: TinyRecursiveReasoningModel_ACTV1Carry, batch: Dict[str, torch.Tensor]) -> Tuple[TinyRecursiveReasoningModel_ACTV1Carry, Dict[str, torch.Tensor]]:
+
+        new_inner_carry, new_steps, new_current_data = self.step_carry(carry, batch)
+
+        # Forward inner model
+        new_inner_carry, outputs = self.inner(new_inner_carry, new_current_data)
+
+
+        with torch.no_grad():
+            # Step
+            new_steps = new_steps + 1
+            is_last_step = new_steps >= self.config.halt_max_steps
+            
+            halted = is_last_step
+
+            # if training, and ACT is enabled
+            if self.training and (self.config.halt_max_steps > 1) or (self.config.eval_with_ACT):
+
+                # Halt signal
+                # NOTE: During evaluation, always use max steps, this is to guarantee the same halting steps inside a batch for batching purposes
+                
+                if self.config.no_ACT_continue:
+                    halted = halted | (outputs["q_halt_logits"] > self.config.ACT_threshold)
+                else:
+                    halted = halted | (outputs["q_halt_logits"] > outputs["q_continue_logits"])
+
+                # Exploration
+                min_halt_steps = (torch.rand_like(outputs["q_halt_logits"]) < self.config.halt_exploration_prob) * torch.randint_like(new_steps, low=2, high=self.config.halt_max_steps + 1)
+                halted = halted & (new_steps >= min_halt_steps)
+
+                if not self.config.no_ACT_continue:
+                    # Compute target Q
+                    # NOTE: No replay buffer and target networks for computing target Q-value.
+                    # As batch_size is large, there're many parallel envs.
+                    # Similar concept as PQN https://arxiv.org/abs/2407.04811
+                    _, _, (next_q_halt_logits, next_q_continue_logits), _, _ = self.inner(new_inner_carry, new_current_data)
+                    outputs["target_q_continue"] = torch.sigmoid(torch.where(is_last_step, next_q_halt_logits, torch.maximum(next_q_halt_logits, next_q_continue_logits)))          
+
+        return TinyRecursiveReasoningModel_ACTV1Carry(new_inner_carry, new_steps, halted, new_current_data), outputs
+
+
+class TinyRecursiveReasoningModel_ACTVClip(nn.Module):
+    """ACT wrapper."""
+
+    def __init__(self, config_dict: dict):
+        super().__init__()
+        self.config = TinyRecursiveReasoningModel_ACTV1Config(**config_dict)
+        self.inner = TinyRecursiveReasoningModel_ACTV1_Inner(self.config)
+
+    @property
+    def puzzle_emb(self):
+        return self.inner.puzzle_emb
+
+    @property
+    def group_emb(self):
+        return self.inner.group_emb
+
+    def initial_carry(self, batch: Dict[str, torch.Tensor]):
+        batch_size = batch["inputs"].shape[0]
+
+        return TinyRecursiveReasoningModel_ACTVClipCarry(
+            inner_carry=self.inner.empty_carry(batch_size),  # Empty is expected, it will be reseted in first pass as all sequences are halted.
+            
+            oldlogits=torch.zeros((batch_size, self.config.seq_len, self.config.vocab_size), dtype=self.inner.forward_dtype),
+            steps=torch.zeros((batch_size, ), dtype=torch.int32),
+            halted=torch.ones((batch_size, ), dtype=torch.bool),  # Default to halted
+            
+            current_data={k: torch.empty_like(v) for k, v in batch.items()}
+        )
+
+    def step_inner_carry(self, carry: TinyRecursiveReasoningModel_ACTVClipCarry, batch: Dict[str, torch.Tensor]):
+        new_inner_carry = self.inner.reset_carry(carry.halted, carry.inner_carry)
+        return new_inner_carry
+
+    def step_current_data(self, carry: TinyRecursiveReasoningModel_ACTVClipCarry, batch: Dict[str, torch.Tensor]):
+        new_current_data = {k: torch.where(carry.halted.view((-1, ) + (1, ) * (batch[k].ndim - 1)), batch[k], v) for k, v in carry.current_data.items()}
+        return new_current_data
+    
+    def step_carry(self, carry: TinyRecursiveReasoningModel_ACTVClipCarry, batch: Dict[str, torch.Tensor]):
+        # Step inner carry
+        new_inner_carry = self.step_inner_carry(carry, batch)
+        
+        # new policy is assumed to start from equal probability 
+        old_logits = torch.where(carry.halted.view((-1, 1, 1)), torch.zeros_like(carry.oldlogits), carry.oldlogits)
+        
+        new_steps = torch.where(carry.halted, 0, carry.steps)
+
+        # Update data, carry (removing halted sequences)        
+        new_current_data = self.step_current_data(carry, batch)
+
+        return new_inner_carry, old_logits, new_steps, new_current_data
+
+
+    def forward(self, carry: TinyRecursiveReasoningModel_ACTVClipCarry, batch: Dict[str, torch.Tensor]) -> Tuple[TinyRecursiveReasoningModel_ACTVClipCarry, Dict[str, torch.Tensor]]:
+
+        new_inner_carry, old_logits, new_steps, new_current_data = self.step_carry(carry, batch)
+
+        # Forward inner model
+        new_inner_carry, outputs = self.inner(new_inner_carry, new_current_data)
+
+
+        with torch.no_grad():
+            # Step
+            new_steps = new_steps + 1
+            is_last_step = new_steps >= self.config.halt_max_steps
+            
+            halted = is_last_step
+
+            # if training, and ACT is enabled
+            if self.training and (self.config.halt_max_steps > 1) or (self.config.eval_with_ACT):
+
+                # Halt signal
+                # NOTE: During evaluation, always use max steps, this is to guarantee the same halting steps inside a batch for batching purposes
+                
+                if self.config.no_ACT_continue:
+                    halted = halted | (outputs["q_halt_logits"] > self.config.ACT_threshold)
+                else:
+                    halted = halted | (outputs["q_halt_logits"] > outputs["q_continue_logits"])
+
+                # Exploration
+                min_halt_steps = (torch.rand_like(outputs["q_halt_logits"]) < self.config.halt_exploration_prob) * torch.randint_like(new_steps, low=2, high=self.config.halt_max_steps + 1)
+                halted = halted & (new_steps >= min_halt_steps)
+
+                if not self.config.no_ACT_continue:
+                    # Compute target Q
+                    # NOTE: No replay buffer and target networks for computing target Q-value.
+                    # As batch_size is large, there're many parallel envs.
+                    # Similar concept as PQN https://arxiv.org/abs/2407.04811
+                    _, _, (next_q_halt_logits, next_q_continue_logits), _, _ = self.inner(new_inner_carry, new_current_data)
+                    outputs["target_q_continue"] = torch.sigmoid(torch.where(is_last_step, next_q_halt_logits, torch.maximum(next_q_halt_logits, next_q_continue_logits)))          
+
+        return TinyRecursiveReasoningModel_ACTVClipCarry(new_inner_carry, old_logits, new_steps, halted, new_current_data), outputs
+
+class TinyRecursiveReasoningModel_ACTGym(TinyRecursiveReasoningModel_ACTV1):
+    """ACT wrapper."""
+
+    def __init__(self, config_dict: dict):
+        super().__init__(config_dict)
+        self.GymConfig = TinyRecursiveReasoningModel_ACTGymConfig(**config_dict)
+
+    def step_inner_carry(self, carry: TinyRecursiveReasoningModel_ACTV1Carry, batch: Dict[str, torch.Tensor]):
+        # check the first step where all halted and steps is 0
+        new_inner_carry = self.inner.reset_carry(carry.halted & (carry.steps > self.GymConfig.replace_halt_threshold), carry.inner_carry)
+        
+        new_inner_carry = self.inner.re_inittialize_carry(carry.halted & (carry.steps <= self.GymConfig.replace_halt_threshold), new_inner_carry, target_latent=self.GymConfig.target_latent_for_reinit, init_std=self.GymConfig.reinit_variance)
+        return new_inner_carry
+
+    def step_carry(self, carry: TinyRecursiveReasoningModel_ACTV1Carry, batch: Dict[str, torch.Tensor]):
+        # Step inner carry
+        new_inner_carry = None
+        if (carry.halted & (carry.steps == 0)).all():
+            new_inner_carry = self.inner.reset_carry(carry.halted, carry.inner_carry)
+        else:
+            new_inner_carry = self.step_inner_carry(carry, batch)
+            carry.halted = carry.halted & (carry.steps > self.GymConfig.replace_halt_threshold)
+
+        new_steps = torch.where(carry.halted, 0, carry.steps)
+
+        # Update data, carry (removing halted sequences)        
+        new_current_data = self.step_current_data(carry, batch)
+
+        return new_inner_carry, new_steps, new_current_data
+
+
+class TinyRecursiveReasoningModel_TRMPPO(TinyRecursiveReasoningModel_ACTV1):
+    """TRM PPO-style wrapper."""
+    
+    def __init__(self, config_dict: dict):
+        super().__init__(config_dict)
+        self.config = TinyRecursiveReasoningModel_ACTV1Config(**config_dict)
+        self.inner = TinyRecursiveReasoningModel_PPOInner(self.config)
+
+    def forward(self, carry: TinyRecursiveReasoningModel_ACTV1Carry, batch: Dict[str, torch.Tensor]) -> Tuple[TinyRecursiveReasoningModel_ACTV1Carry, Dict[str, torch.Tensor]]:
+
+        new_inner_carry, new_steps, new_current_data = self.step_carry(carry, batch)
+
+        # Forward inner model
+        new_inner_carry, outputs = self.inner(new_inner_carry, new_current_data)
+
+
+        with torch.no_grad():
+            # Step
+            new_steps = new_steps + 1
+            is_last_step = new_steps >= self.config.halt_max_steps
+            
+            halted = is_last_step
+
+        return TinyRecursiveReasoningModel_ACTV1Carry(new_inner_carry, new_steps, halted, new_current_data), outputs