Delete sudoku/trm.py

sudoku/trm.py

@@ -1,645 +0,0 @@
-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