diff --git a/.gitattributes b/.gitattributes
index a6344aac8c09253b3b630fb776ae94478aa0275b..8f79f511f284a0bd2cb6cc4afc8876b9d30a6196 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -33,3 +33,63 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_10416 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_11718 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_1302 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_13020 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_14322 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_15624 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_16926 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_18228 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_19530 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_20832 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_22134 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_23436 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_24738 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_2604 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_26040 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_27342 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_28644 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_29946 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_31248 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_32550 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_33852 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_35154 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_36456 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_37758 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_3906 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_39060 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_40362 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_41664 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_42966 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_44268 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_45570 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_46872 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_48174 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_49476 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_50778 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_5208 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_52080 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_6510 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_7812 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/augmented-hrm/step_9114 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/original-hrm/step_10416 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/original-hrm/step_13020 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/original-hrm/step_15624 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/original-hrm/step_18228 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/original-hrm/step_20832 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/original-hrm/step_23436 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/original-hrm/step_2604 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/original-hrm/step_26040 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/original-hrm/step_28644 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/original-hrm/step_31248 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/original-hrm/step_33852 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/original-hrm/step_36456 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/original-hrm/step_39060 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/original-hrm/step_41664 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/original-hrm/step_44268 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/original-hrm/step_46872 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/original-hrm/step_49476 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/original-hrm/step_5208 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/original-hrm/step_52080 filter=lfs diff=lfs merge=lfs -text
+sudoku-extreme/original-hrm/step_7812 filter=lfs diff=lfs merge=lfs -text
diff --git a/sudoku-extreme/augmented-hrm/all_config.yaml b/sudoku-extreme/augmented-hrm/all_config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..357d4c30799cbe9b27ebdcea3a46b10e3d640111
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/all_config.yaml
@@ -0,0 +1,35 @@
+arch:
+  H_cycles: 2
+  H_layers: 4
+  L_cycles: 2
+  L_layers: 4
+  expansion: 4
+  halt_exploration_prob: 0.1
+  halt_max_steps: 16
+  hidden_size: 512
+  loss:
+    loss_type: stablemax_cross_entropy
+    name: losses@ACTLossHead
+  name: hrm.hrm_act_v1@HierarchicalReasoningModel_ACTV1
+  num_heads: 8
+  pos_encodings: rope
+  puzzle_emb_ndim: 512
+beta1: 0.9
+beta2: 0.95
+checkpoint_every_eval: true
+checkpoint_path: checkpoints/Sudoku-extreme-1k-aug-1000-hint ACT-torch/HierarchicalReasoningModel_ACTV1
+  hopeful-quetzal
+data_path: ../../dataset/data/sudoku-extreme-1k-aug-1000-hint
+epochs: 40000
+eval_interval: 1000
+eval_save_outputs: []
+global_batch_size: 768
+lr: 0.0001
+lr_min_ratio: 1.0
+lr_warmup_steps: 2000
+project_name: Sudoku-extreme-1k-aug-1000-hint ACT-torch
+puzzle_emb_lr: 0.0001
+puzzle_emb_weight_decay: 1.0
+run_name: HierarchicalReasoningModel_ACTV1 hopeful-quetzal
+seed: 0
+weight_decay: 1.0
diff --git a/sudoku-extreme/augmented-hrm/hrm_act_v1.py b/sudoku-extreme/augmented-hrm/hrm_act_v1.py
new file mode 100644
index 0000000000000000000000000000000000000000..e0a253ca7c995d4d18449e27ac9f242dbedad7fa
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/hrm_act_v1.py
@@ -0,0 +1,311 @@
+from typing import Tuple, List, Dict, Optional
+from dataclasses import dataclass
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from pydantic import BaseModel
+
+from models.common import trunc_normal_init_
+from models.layers import rms_norm, SwiGLU, Attention, RotaryEmbedding, CosSin, CastedEmbedding, CastedLinear
+from models.sparse_embedding import CastedSparseEmbedding
+
+
+@dataclass
+class HierarchicalReasoningModel_ACTV1InnerCarry:
+    z_H: torch.Tensor
+    z_L: torch.Tensor
+
+
+@dataclass
+class HierarchicalReasoningModel_ACTV1Carry:
+    inner_carry: HierarchicalReasoningModel_ACTV1InnerCarry
+    
+    steps: torch.Tensor
+    halted: torch.Tensor
+    
+    current_data: Dict[str, torch.Tensor]
+
+
+class HierarchicalReasoningModel_ACTV1Config(BaseModel):
+    batch_size: int
+    seq_len: int
+    puzzle_emb_ndim: int = 0
+    num_puzzle_identifiers: int
+    vocab_size: int
+
+    H_cycles: int
+    L_cycles: int
+
+    H_layers: int
+    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
+
+    forward_dtype: str = "bfloat16"
+
+
+class HierarchicalReasoningModel_ACTV1Block(nn.Module):
+    def __init__(self, config: HierarchicalReasoningModel_ACTV1Config) -> None:
+        super().__init__()
+
+        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:
+        # Post Norm
+        # 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
+        hidden_states = rms_norm(hidden_states + self.mlp(hidden_states), variance_epsilon=self.norm_eps)
+        return hidden_states
+
+
+class HierarchicalReasoningModel_ACTV1ReasoningModule(nn.Module):
+    def __init__(self, layers: List[HierarchicalReasoningModel_ACTV1Block]):
+        super().__init__()
+
+        self.layers = torch.nn.ModuleList(layers)
+
+    def forward(self, hidden_states: torch.Tensor, input_injection: torch.Tensor, **kwargs) -> torch.Tensor:
+        # Input injection (add)
+        hidden_states = hidden_states + input_injection
+        # Layers
+        for layer in self.layers:
+            hidden_states = layer(hidden_states=hidden_states, **kwargs)
+
+        return hidden_states
+
+
+class HierarchicalReasoningModel_ACTV1_Inner(nn.Module):
+    def __init__(self, config: HierarchicalReasoningModel_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)
+
+        self.puzzle_emb_len = -(self.config.puzzle_emb_ndim // -self.config.hidden_size)  # ceil div
+        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)
+
+        # 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:
+            raise NotImplementedError()
+
+        # Reasoning Layers
+        self.H_level = HierarchicalReasoningModel_ACTV1ReasoningModule(layers=[HierarchicalReasoningModel_ACTV1Block(self.config) for _i in range(self.config.H_layers)])
+        self.L_level = HierarchicalReasoningModel_ACTV1ReasoningModule(layers=[HierarchicalReasoningModel_ACTV1Block(self.config) for _i in range(self.config.L_layers)])
+        
+        # Initial states
+        self.H_init = nn.Buffer(trunc_normal_init_(torch.empty(self.config.hidden_size, dtype=self.forward_dtype), std=1), persistent=True)
+        self.L_init = nn.Buffer(trunc_normal_init_(torch.empty(self.config.hidden_size, dtype=self.forward_dtype), std=1), persistent=True)
+
+        # 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
+
+    def _input_embeddings(self, input: torch.Tensor, puzzle_identifiers: 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)
+            
+            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 HierarchicalReasoningModel_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: HierarchicalReasoningModel_ACTV1InnerCarry, use_default=True):
+        if use_default:
+            return HierarchicalReasoningModel_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),
+            )
+        else:
+            H_tmp = trunc_normal_init_(torch.empty(self.config.hidden_size, dtype=self.forward_dtype), std=2).to("cuda")+self.H_init
+            L_tmp = trunc_normal_init_(torch.empty(self.config.hidden_size, dtype=self.forward_dtype), std=2).to("cuda")+self.L_init
+            return HierarchicalReasoningModel_ACTV1InnerCarry(
+                z_H=torch.where(reset_flag.view(-1, 1, 1), H_tmp, carry.z_H),
+                z_L=torch.where(reset_flag.view(-1, 1, 1), L_tmp, carry.z_L)
+            )
+
+
+    def forward(self, carry: HierarchicalReasoningModel_ACTV1InnerCarry, batch: Dict[str, torch.Tensor], require_trace=False):
+    #  -> Tuple[HierarchicalReasoningModel_ACTV1InnerCarry, torch.Tensor, Tuple[torch.Tensor, 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"])
+
+        z_H_trace = []
+
+        # Forward iterations
+        with torch.no_grad():
+            z_H, z_L = carry.z_H, carry.z_L
+
+            for _H_step in range(self.config.H_cycles):
+                for _L_step in range(self.config.L_cycles):
+                    if not ((_H_step == self.config.H_cycles - 1) and (_L_step == self.config.L_cycles - 1)):
+                        z_L = self.L_level(z_L, z_H + input_embeddings, **seq_info)
+
+                if not (_H_step == self.config.H_cycles - 1):
+                    z_H = self.H_level(z_H, z_L, **seq_info)
+                    if require_trace:
+                        z_H_trace.append(z_H.detach().cpu().clone())
+
+        assert not z_H.requires_grad and not z_L.requires_grad
+
+        # 1-step grad
+        z_L = self.L_level(z_L, z_H + input_embeddings, **seq_info)
+        z_H = self.H_level(z_H, z_L, **seq_info)
+
+        if require_trace:
+            z_H_trace.append(z_H.detach().cpu().clone())
+
+        # LM Outputs
+        new_carry = HierarchicalReasoningModel_ACTV1InnerCarry(z_H=z_H.detach(), z_L=z_L.detach())  # New carry no grad
+        output = self.lm_head(z_H)[:, self.puzzle_emb_len:]
+
+        # Q head
+        q_logits = self.q_head(z_H[:, 0]).to(torch.float32)
+        
+        if require_trace:
+            return z_H_trace, new_carry, output, (q_logits[..., 0], q_logits[..., 1])
+        else:
+            return new_carry, output, (q_logits[..., 0], q_logits[..., 1])
+
+
+class HierarchicalReasoningModel_ACTV1(nn.Module):
+    """ACT wrapper."""
+
+    def __init__(self, config_dict: dict):
+        super().__init__()
+        self.config = HierarchicalReasoningModel_ACTV1Config(**config_dict)
+        self.inner = HierarchicalReasoningModel_ACTV1_Inner(self.config)
+
+    @property
+    def puzzle_emb(self):
+        return self.inner.puzzle_emb
+
+    def initial_carry(self, batch: Dict[str, torch.Tensor]):
+        batch_size = batch["inputs"].shape[0]
+
+        return HierarchicalReasoningModel_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 forward(self, carry: HierarchicalReasoningModel_ACTV1Carry, batch: Dict[str, torch.Tensor], require_trace=False):
+        # -> Tuple[HierarchicalReasoningModel_ACTV1Carry, Dict[str, torch.Tensor], torch.Tensor]:
+        # Update data, carry (removing halted sequences)
+        new_inner_carry = self.inner.reset_carry(carry.halted, carry.inner_carry)
+        
+        new_steps = torch.where(carry.halted, 0, carry.steps)
+
+        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()}
+
+        # Forward inner model
+        if require_trace:
+            z_H_trace, new_inner_carry, logits, (q_halt_logits, q_continue_logits) = self.inner(new_inner_carry, new_current_data, require_trace=require_trace)
+        else:
+            new_inner_carry, logits, (q_halt_logits, q_continue_logits) = self.inner(new_inner_carry, new_current_data)
+
+        outputs = {
+            "logits": logits,
+            "q_halt_logits": q_halt_logits,
+            "q_continue_logits": q_continue_logits
+        }
+        
+        with torch.no_grad():
+            # Step
+            new_steps = new_steps + 1
+            is_last_step = new_steps >= self.config.halt_max_steps
+            # is_last_step = new_steps >= 32
+            
+            halted = is_last_step
+
+            # if training, and ACT is enabled
+            if self.training and (self.config.halt_max_steps > 1):
+                # Halt signal
+                # NOTE: During evaluation, always use max steps, this is to guarantee the same halting steps inside a batch for batching purposes
+                halted = halted | (q_halt_logits > q_continue_logits)
+
+                # Exploration
+                min_halt_steps = (torch.rand_like(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)
+
+                # 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)[-1]
+                
+                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)))
+
+        if require_trace:
+            return HierarchicalReasoningModel_ACTV1Carry(new_inner_carry, new_steps, halted, new_current_data), outputs, new_steps, (q_halt_logits > q_continue_logits), z_H_trace
+        else:
+            return HierarchicalReasoningModel_ACTV1Carry(new_inner_carry, new_steps, halted, new_current_data), outputs, new_steps, (q_halt_logits > q_continue_logits)
diff --git a/sudoku-extreme/augmented-hrm/losses.py b/sudoku-extreme/augmented-hrm/losses.py
new file mode 100644
index 0000000000000000000000000000000000000000..45bea6882babe8b0310e1d733ecb7db672004b95
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/losses.py
@@ -0,0 +1,115 @@
+from typing import Any, Tuple, Dict, Sequence, Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+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):
+    logprobs = log_stablemax(logits.to(torch.float64), dim=-1)
+
+    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)
+
+
+class ACTLossHead(nn.Module):
+    def __init__(self, model: nn.Module, loss_type: str):
+        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
+
+    def forward(
+        self,
+        return_keys: Sequence[str],
+        require_trace=False,
+        # 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
+        if require_trace:
+            new_carry, outputs, steps, act_halt, z_H_trace = self.model(**model_kwargs, require_trace=require_trace)
+        else:
+            new_carry, outputs, steps, act_halt = self.model(**model_kwargs)
+        labels = new_carry.current_data["labels"]
+        alpha = 1.0
+        ds_mask = (alpha ** (16-steps.detach())).unsqueeze(dim=1)
+        # print(ds_mask.shape)
+
+        # Correctness
+        with torch.no_grad():
+            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)
+            valid_metrics = (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
+        # FIXME: Assuming the batch is always full
+        loss = self.loss_fn(outputs["logits"], labels, ignore_index=IGNORE_LABEL_ID)
+        # print(loss.shape)
+        lm_loss = (ds_mask * loss / loss_divisor).sum()
+        q_halt_loss = F.binary_cross_entropy_with_logits(outputs["q_halt_logits"], seq_is_correct.to(outputs["q_halt_logits"].dtype), reduction="sum")
+
+        metrics.update({
+            "lm_loss": lm_loss.detach(),
+            "q_halt_loss": q_halt_loss.detach(),
+        })
+
+        # Q continue (bootstrapping target loss)
+        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}
+
+        if require_trace:
+            return z_H_trace, new_carry, lm_loss + 0.5 * (q_halt_loss + q_continue_loss), metrics, detached_outputs, new_carry.halted.all(), new_carry.halted & act_halt
+        else:
+            return new_carry, lm_loss + 0.5 * (q_halt_loss + q_continue_loss), metrics, detached_outputs, new_carry.halted.all(), new_carry.halted & act_halt
diff --git a/sudoku-extreme/augmented-hrm/step_10416 b/sudoku-extreme/augmented-hrm/step_10416
new file mode 100644
index 0000000000000000000000000000000000000000..eff625e1287af333b162620e6c8b585e9740e4bb
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_10416
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c58c3f870e01c7b9ea71acd128383eb1121ead36b039ddc527ba99326df4a27
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_11718 b/sudoku-extreme/augmented-hrm/step_11718
new file mode 100644
index 0000000000000000000000000000000000000000..e4612c8f9c6a8f3363787590b99a8a8ab24eb4b1
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_11718
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6b8a17f320bcf07ec9d4de1a8d0700906ac6272bc21645beb495f6e63130e40f
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_1302 b/sudoku-extreme/augmented-hrm/step_1302
new file mode 100644
index 0000000000000000000000000000000000000000..f14415d493aa7e773a01688470b8f255a0b25617
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_1302
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6b4184cfd2db078c6a51160cb52eb501b5055d4570817a4e35015069afcf2163
+size 109124296
diff --git a/sudoku-extreme/augmented-hrm/step_13020 b/sudoku-extreme/augmented-hrm/step_13020
new file mode 100644
index 0000000000000000000000000000000000000000..7327aac65438a3a30da0cbdc6b300998ee913ffa
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_13020
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d897d505722435b39b27df0371df24db25e65b16aa726027741fad22551eeace
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_14322 b/sudoku-extreme/augmented-hrm/step_14322
new file mode 100644
index 0000000000000000000000000000000000000000..d439734369395da9a73adde9c47d967c92c682df
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_14322
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:27f34b2a089dc0dfb798ac8a803bb6881e04b812d73f4ccd4caa1403106e635d
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_15624 b/sudoku-extreme/augmented-hrm/step_15624
new file mode 100644
index 0000000000000000000000000000000000000000..1c3ea50e5991305f13d30b216d817c5d2aca4d3b
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_15624
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:055c2210483f46979a4f3dba036f5849376c4ad94306e238bc20c6bdba9385d5
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_16926 b/sudoku-extreme/augmented-hrm/step_16926
new file mode 100644
index 0000000000000000000000000000000000000000..60249bdce2655020668246a50e573ed3d1ca0a81
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_16926
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0f6c65467f330309d837724f11ef5be112ab15c59153a054d0620d6ce038bd93
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_18228 b/sudoku-extreme/augmented-hrm/step_18228
new file mode 100644
index 0000000000000000000000000000000000000000..82269f32e0b83bc35594c473d8984bcd6636a59f
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_18228
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1f09e18eda6656bbf76a1731f1dd83d4ac96cc0289bb6dc61f836d87e2b9744a
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_19530 b/sudoku-extreme/augmented-hrm/step_19530
new file mode 100644
index 0000000000000000000000000000000000000000..a4a2a0ef79d1ccd0bd3ff5862262cbd9ec24ead5
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_19530
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:69af9c6808313c91014c3f493a581333c4241288780040f74e027f8011c73da1
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_20832 b/sudoku-extreme/augmented-hrm/step_20832
new file mode 100644
index 0000000000000000000000000000000000000000..ad29583b30f6f338c6728319a3915a417e596d53
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_20832
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:89c400212dbcdf4ad9aad1b8bf2da77a3e0f5f72a97e9f01eebc00bd5c50987f
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_22134 b/sudoku-extreme/augmented-hrm/step_22134
new file mode 100644
index 0000000000000000000000000000000000000000..230602c6759986f216bb4f962ac47c2f757ec3b0
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_22134
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9b6664aac34744f87463250d379d2d3fb75eb35229e609b640ea5e744a275f09
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_23436 b/sudoku-extreme/augmented-hrm/step_23436
new file mode 100644
index 0000000000000000000000000000000000000000..663539c685e6692e753dfabb7ab92225dc6f5e6f
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_23436
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:90015d8bf2a1164e7b76445e187b2944343434d40eab16c7adbad65cbbe183c1
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_24738 b/sudoku-extreme/augmented-hrm/step_24738
new file mode 100644
index 0000000000000000000000000000000000000000..ef7eeb1daa1f8b6713b637bc36386bb770b13c62
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_24738
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:194ef810eeed7c3c6ab47a6dcecd005973d94f64cebaba28b735149672e57bc9
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_2604 b/sudoku-extreme/augmented-hrm/step_2604
new file mode 100644
index 0000000000000000000000000000000000000000..389d289bbc0fd9ab4bc93e19915581a691f69eb1
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_2604
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bc3077865661b89a41f15b9e1d3d7500278eab43681a762557328c3965c9613e
+size 109124296
diff --git a/sudoku-extreme/augmented-hrm/step_26040 b/sudoku-extreme/augmented-hrm/step_26040
new file mode 100644
index 0000000000000000000000000000000000000000..27ff8af05d831df77f19be34bc2f5c577e5862a0
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_26040
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:822db0d53159ea6e3609e0231145b11644661f137ce154e69932f2799c6fc218
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_27342 b/sudoku-extreme/augmented-hrm/step_27342
new file mode 100644
index 0000000000000000000000000000000000000000..99ec20ab2d454ce18d8c0559f490939a225053af
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_27342
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:21d01ac7893b791d776d718d312258ba163e34cb83f79681a3199cb16da3cc7b
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_28644 b/sudoku-extreme/augmented-hrm/step_28644
new file mode 100644
index 0000000000000000000000000000000000000000..8b4f5713996e83b36847761c541919c3f58b45b2
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_28644
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0a4a79b8c9257e519821e9fc2256a2bbee2259a54577f1feda15a5daa0f419af
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_29946 b/sudoku-extreme/augmented-hrm/step_29946
new file mode 100644
index 0000000000000000000000000000000000000000..f8217608ac031c5285b4b7a233966550de3ac887
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_29946
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:62691d4a98c1dbbcfc9df55f3228bf28d9302b97bcfa248bd877227aea5afa04
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_31248 b/sudoku-extreme/augmented-hrm/step_31248
new file mode 100644
index 0000000000000000000000000000000000000000..5d625dc6ba390a0ad8f1f3cc68575e87b356f647
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_31248
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:907f28224104af1123750c968e6bff566871a5fa3dd0ec2602622ca0cdbb5431
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_32550 b/sudoku-extreme/augmented-hrm/step_32550
new file mode 100644
index 0000000000000000000000000000000000000000..6363cc286377c7c7866b0e25ff2ed8011cf95fa4
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_32550
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd2146110c8af9a8f59144c9e2f1138aa8c8b354e67e4939e2adeb394080a54e
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_33852 b/sudoku-extreme/augmented-hrm/step_33852
new file mode 100644
index 0000000000000000000000000000000000000000..15844154e30d634c80394fa5237dd410e5367bc6
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_33852
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fbd7dfd484acfa11daa48a0e7df5b97a402dd1891e1a770953dce25ba3212311
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_35154 b/sudoku-extreme/augmented-hrm/step_35154
new file mode 100644
index 0000000000000000000000000000000000000000..cb41de869b6358fe4df5801c20af21b243261670
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_35154
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e137a7095ed637f383da29fa51d548d7f8381dd852f15f98bd918c2754cb3472
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_36456 b/sudoku-extreme/augmented-hrm/step_36456
new file mode 100644
index 0000000000000000000000000000000000000000..102fdb40ec541e99c13031c7abc1868d36dcc7e2
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_36456
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ddb9bf9d42109d1c13cd6a39c4e86a19555b637c38bf5b992deb6c10bffe4bde
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_37758 b/sudoku-extreme/augmented-hrm/step_37758
new file mode 100644
index 0000000000000000000000000000000000000000..a5aabf4acb0d84d39b5b50e16522340624629dac
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_37758
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:70e86d638e876d29fdd4b7d07d68460dac779e402a7fda4da137edb73c5b3f79
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_3906 b/sudoku-extreme/augmented-hrm/step_3906
new file mode 100644
index 0000000000000000000000000000000000000000..52e4a0f8d7211000048cff1d16ae2a01060167fa
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_3906
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2b3bcf655ca6bda5eeababb44b736d4da76a29fbfafc08378248341177b4f29d
+size 109124296
diff --git a/sudoku-extreme/augmented-hrm/step_39060 b/sudoku-extreme/augmented-hrm/step_39060
new file mode 100644
index 0000000000000000000000000000000000000000..6110e8991255bf46d85a59e278508d9258ec72a7
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_39060
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c2b141ae8fbb6333801e44ea3033e381723a01c513ae5115d13fb1606ea017cb
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_40362 b/sudoku-extreme/augmented-hrm/step_40362
new file mode 100644
index 0000000000000000000000000000000000000000..0af619a29bc7cea103469ea7ce11a6e8011619d2
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_40362
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:85be48ffd29247f0b2955596e40bce1c4d3f426db3ef2f14aaada46cffe64e47
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_41664 b/sudoku-extreme/augmented-hrm/step_41664
new file mode 100644
index 0000000000000000000000000000000000000000..de544565d55a1f08351e8460e69bd0aa108bb11c
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_41664
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:331ca95e1e69985de2b9585fa3d908084b76d6595330d6701dc6cc722bf5cf5f
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_42966 b/sudoku-extreme/augmented-hrm/step_42966
new file mode 100644
index 0000000000000000000000000000000000000000..8a6ead300b65ed15ce782801ea110484a607fc14
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_42966
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:763d083485336726a4d8a8fc96d48efa0f361152f3987b29b00521ddedde7379
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_44268 b/sudoku-extreme/augmented-hrm/step_44268
new file mode 100644
index 0000000000000000000000000000000000000000..31730ef40d75ac48c315cb7c2b84521b02127dfc
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_44268
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5512a6b339ba9e4f82c32d7715a18b2efb00dbbbe30752e4385be226cb896291
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_45570 b/sudoku-extreme/augmented-hrm/step_45570
new file mode 100644
index 0000000000000000000000000000000000000000..1b636feda14e0f50f6ec2c1718fcf67b48f0e2dd
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_45570
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08fab5e91dbac51dd230c38ff9722f54bab1bcba572eeb2680b5a3dea38c37f3
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_46872 b/sudoku-extreme/augmented-hrm/step_46872
new file mode 100644
index 0000000000000000000000000000000000000000..d3716674988b4a1fe464308dd9008c4e435d9504
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_46872
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c38b064cb860b4010d67bcd436a65074db539eb7650bdabeff6857bbbe101909
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_48174 b/sudoku-extreme/augmented-hrm/step_48174
new file mode 100644
index 0000000000000000000000000000000000000000..da458656ea81acd6447a9e1efa8026f862ab2d31
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_48174
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ad5637a7444577632c9d940ef28ec6726472bddf0f24885b09697372351711f4
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_49476 b/sudoku-extreme/augmented-hrm/step_49476
new file mode 100644
index 0000000000000000000000000000000000000000..b2cda1e1f1b3c5046084225a3bdfa789775f9db7
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_49476
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:648b7ae9f110efedff7c9da31afd90953ce1344a08154de5866e4efcfc13d9c8
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_50778 b/sudoku-extreme/augmented-hrm/step_50778
new file mode 100644
index 0000000000000000000000000000000000000000..da0a315b3a6d43f9216dddb11f6600a0c9fb5e9a
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_50778
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:60da15d0545b8eaa638c5ed888a7e3a12116a5c575d55b1fb22abb03f2183712
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_5208 b/sudoku-extreme/augmented-hrm/step_5208
new file mode 100644
index 0000000000000000000000000000000000000000..77a6845bab9b6f9d7e8e6f48aac3a358b151c229
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_5208
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:56a4b69e82b45e38c0c5aaa6dcf4b28109995f6283d2880b4f15e982c713d1e4
+size 109124296
diff --git a/sudoku-extreme/augmented-hrm/step_52080 b/sudoku-extreme/augmented-hrm/step_52080
new file mode 100644
index 0000000000000000000000000000000000000000..56cfe21366115073d486d73ab0fc2eda1191598d
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_52080
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b83423bb0b263399a72c6259c3527993a427a00792e9efd3e0b0d87c21b6b7d6
+size 109124341
diff --git a/sudoku-extreme/augmented-hrm/step_6510 b/sudoku-extreme/augmented-hrm/step_6510
new file mode 100644
index 0000000000000000000000000000000000000000..f3b79c81fdd2f948d297478bc467b71b302073da
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_6510
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aba4bd3427b9b524df4fb4d45f757efef72a2b5c40c7ae383c67088f3fc0d0eb
+size 109124296
diff --git a/sudoku-extreme/augmented-hrm/step_7812 b/sudoku-extreme/augmented-hrm/step_7812
new file mode 100644
index 0000000000000000000000000000000000000000..aee302ff1fcf35e3f00ca8f8647b9d60ffeef44e
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_7812
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4783d7ce493907959a2d82d719103092b69c0f295695f62596dc19368c7ecc76
+size 109124296
diff --git a/sudoku-extreme/augmented-hrm/step_9114 b/sudoku-extreme/augmented-hrm/step_9114
new file mode 100644
index 0000000000000000000000000000000000000000..df93257daf9a31a91987b6e8f4c611fcbdf0ce22
--- /dev/null
+++ b/sudoku-extreme/augmented-hrm/step_9114
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1d0bae068e1c11e7550f4c59f918abe22a3986b9df995d5282d61541c3633acc
+size 109124296
diff --git a/sudoku-extreme/original-hrm/all_config.yaml b/sudoku-extreme/original-hrm/all_config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..bfd2854f767da9c5509ac1c58371ecb441b33a19
--- /dev/null
+++ b/sudoku-extreme/original-hrm/all_config.yaml
@@ -0,0 +1,35 @@
+arch:
+  H_cycles: 2
+  H_layers: 4
+  L_cycles: 2
+  L_layers: 4
+  expansion: 4
+  halt_exploration_prob: 0.1
+  halt_max_steps: 16
+  hidden_size: 512
+  loss:
+    loss_type: stablemax_cross_entropy
+    name: losses@ACTLossHead
+  name: hrm.hrm_act_v1@HierarchicalReasoningModel_ACTV1
+  num_heads: 8
+  pos_encodings: rope
+  puzzle_emb_ndim: 512
+beta1: 0.9
+beta2: 0.95
+checkpoint_every_eval: true
+checkpoint_path: checkpoints/Sudoku-extreme-1k-aug-1000 ACT-torch/HierarchicalReasoningModel_ACTV1
+  liberal-bee
+data_path: ../../../dataset/data/sudoku-extreme-1k-aug-1000
+epochs: 20000
+eval_interval: 1000
+eval_save_outputs: []
+global_batch_size: 384
+lr: 7.0e-05
+lr_min_ratio: 1.0
+lr_warmup_steps: 2000
+project_name: Sudoku-extreme-1k-aug-1000 ACT-torch
+puzzle_emb_lr: 7.0e-05
+puzzle_emb_weight_decay: 1.0
+run_name: HierarchicalReasoningModel_ACTV1 liberal-bee
+seed: 0
+weight_decay: 1.0
diff --git a/sudoku-extreme/original-hrm/hrm_act_v1.py b/sudoku-extreme/original-hrm/hrm_act_v1.py
new file mode 100644
index 0000000000000000000000000000000000000000..e91c7d1a047b9f9735fa8253daa4b8effa23f011
--- /dev/null
+++ b/sudoku-extreme/original-hrm/hrm_act_v1.py
@@ -0,0 +1,283 @@
+from typing import Tuple, List, Dict, Optional
+from dataclasses import dataclass
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from pydantic import BaseModel
+
+from models.common import trunc_normal_init_
+from models.layers import rms_norm, SwiGLU, Attention, RotaryEmbedding, CosSin, CastedEmbedding, CastedLinear
+from models.sparse_embedding import CastedSparseEmbedding
+
+
+@dataclass
+class HierarchicalReasoningModel_ACTV1InnerCarry:
+    z_H: torch.Tensor
+    z_L: torch.Tensor
+
+
+@dataclass
+class HierarchicalReasoningModel_ACTV1Carry:
+    inner_carry: HierarchicalReasoningModel_ACTV1InnerCarry
+    
+    steps: torch.Tensor
+    halted: torch.Tensor
+    
+    current_data: Dict[str, torch.Tensor]
+
+
+class HierarchicalReasoningModel_ACTV1Config(BaseModel):
+    batch_size: int
+    seq_len: int
+    puzzle_emb_ndim: int = 0
+    num_puzzle_identifiers: int
+    vocab_size: int
+
+    H_cycles: int
+    L_cycles: int
+
+    H_layers: int
+    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
+
+    forward_dtype: str = "bfloat16"
+
+
+class HierarchicalReasoningModel_ACTV1Block(nn.Module):
+    def __init__(self, config: HierarchicalReasoningModel_ACTV1Config) -> None:
+        super().__init__()
+
+        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:
+        # Post Norm
+        # 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
+        hidden_states = rms_norm(hidden_states + self.mlp(hidden_states), variance_epsilon=self.norm_eps)
+        return hidden_states
+
+
+class HierarchicalReasoningModel_ACTV1ReasoningModule(nn.Module):
+    def __init__(self, layers: List[HierarchicalReasoningModel_ACTV1Block]):
+        super().__init__()
+
+        self.layers = torch.nn.ModuleList(layers)
+
+    def forward(self, hidden_states: torch.Tensor, input_injection: torch.Tensor, **kwargs) -> torch.Tensor:
+        # Input injection (add)
+        hidden_states = hidden_states + input_injection
+        # Layers
+        for layer in self.layers:
+            hidden_states = layer(hidden_states=hidden_states, **kwargs)
+
+        return hidden_states
+
+
+class HierarchicalReasoningModel_ACTV1_Inner(nn.Module):
+    def __init__(self, config: HierarchicalReasoningModel_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)
+
+        self.puzzle_emb_len = -(self.config.puzzle_emb_ndim // -self.config.hidden_size)  # ceil div
+        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)
+
+        # 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:
+            raise NotImplementedError()
+
+        # Reasoning Layers
+        self.H_level = HierarchicalReasoningModel_ACTV1ReasoningModule(layers=[HierarchicalReasoningModel_ACTV1Block(self.config) for _i in range(self.config.H_layers)])
+        self.L_level = HierarchicalReasoningModel_ACTV1ReasoningModule(layers=[HierarchicalReasoningModel_ACTV1Block(self.config) for _i in range(self.config.L_layers)])
+        
+        # Initial states
+        self.H_init = nn.Buffer(trunc_normal_init_(torch.empty(self.config.hidden_size, dtype=self.forward_dtype), std=1), persistent=True)
+        self.L_init = nn.Buffer(trunc_normal_init_(torch.empty(self.config.hidden_size, dtype=self.forward_dtype), std=1), persistent=True)
+
+        # 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
+
+    def _input_embeddings(self, input: torch.Tensor, puzzle_identifiers: 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)
+            
+            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 HierarchicalReasoningModel_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: HierarchicalReasoningModel_ACTV1InnerCarry):
+        return HierarchicalReasoningModel_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 forward(self, carry: HierarchicalReasoningModel_ACTV1InnerCarry, batch: Dict[str, torch.Tensor]) -> Tuple[HierarchicalReasoningModel_ACTV1InnerCarry, torch.Tensor, Tuple[torch.Tensor, 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"])
+
+        # Forward iterations
+        with torch.no_grad():
+            z_H, z_L = carry.z_H, carry.z_L
+
+            for _H_step in range(self.config.H_cycles):
+                for _L_step in range(self.config.L_cycles):
+                    if not ((_H_step == self.config.H_cycles - 1) and (_L_step == self.config.L_cycles - 1)):
+                        z_L = self.L_level(z_L, z_H + input_embeddings, **seq_info)
+
+                if not (_H_step == self.config.H_cycles - 1):
+                    z_H = self.H_level(z_H, z_L, **seq_info)
+
+        assert not z_H.requires_grad and not z_L.requires_grad
+
+        # 1-step grad
+        z_L = self.L_level(z_L, z_H + input_embeddings, **seq_info)
+        z_H = self.H_level(z_H, z_L, **seq_info)
+
+        # LM Outputs
+        new_carry = HierarchicalReasoningModel_ACTV1InnerCarry(z_H=z_H.detach(), z_L=z_L.detach())  # New carry no grad
+        output = self.lm_head(z_H)[:, self.puzzle_emb_len:]
+
+        # Q head
+        q_logits = self.q_head(z_H[:, 0]).to(torch.float32)
+        
+        return new_carry, output, (q_logits[..., 0], q_logits[..., 1])
+
+
+class HierarchicalReasoningModel_ACTV1(nn.Module):
+    """ACT wrapper."""
+
+    def __init__(self, config_dict: dict):
+        super().__init__()
+        self.config = HierarchicalReasoningModel_ACTV1Config(**config_dict)
+        self.inner = HierarchicalReasoningModel_ACTV1_Inner(self.config)
+
+    @property
+    def puzzle_emb(self):
+        return self.inner.puzzle_emb
+
+    def initial_carry(self, batch: Dict[str, torch.Tensor]):
+        batch_size = batch["inputs"].shape[0]
+
+        return HierarchicalReasoningModel_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 forward(self, carry: HierarchicalReasoningModel_ACTV1Carry, batch: Dict[str, torch.Tensor]) -> Tuple[HierarchicalReasoningModel_ACTV1Carry, Dict[str, torch.Tensor]]:
+        # Update data, carry (removing halted sequences)
+        new_inner_carry = self.inner.reset_carry(carry.halted, carry.inner_carry)
+        
+        new_steps = torch.where(carry.halted, 0, carry.steps)
+
+        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()}
+
+        # Forward inner model
+        new_inner_carry, logits, (q_halt_logits, q_continue_logits) = self.inner(new_inner_carry, new_current_data)
+
+        outputs = {
+            "logits": logits,
+            "q_halt_logits": q_halt_logits,
+            "q_continue_logits": q_continue_logits
+        }
+        
+        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):
+                # Halt signal
+                # NOTE: During evaluation, always use max steps, this is to guarantee the same halting steps inside a batch for batching purposes
+                halted = halted | (q_halt_logits > q_continue_logits)
+
+                # Exploration
+                min_halt_steps = (torch.rand_like(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)
+
+                # 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)[-1]
+                
+                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 HierarchicalReasoningModel_ACTV1Carry(new_inner_carry, new_steps, halted, new_current_data), outputs
diff --git a/sudoku-extreme/original-hrm/losses.py b/sudoku-extreme/original-hrm/losses.py
new file mode 100644
index 0000000000000000000000000000000000000000..b3118e728365cc2824646a156a35d79f6dcca9c1
--- /dev/null
+++ b/sudoku-extreme/original-hrm/losses.py
@@ -0,0 +1,101 @@
+from typing import Any, Tuple, Dict, Sequence, Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+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):
+    logprobs = log_stablemax(logits.to(torch.float64), dim=-1)
+
+    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)
+
+
+class ACTLossHead(nn.Module):
+    def __init__(self, model: nn.Module, loss_type: str):
+        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
+
+    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"]
+
+        # Correctness
+        with torch.no_grad():
+            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
+        # FIXME: Assuming the batch is always full
+        lm_loss = (self.loss_fn(outputs["logits"], labels, ignore_index=IGNORE_LABEL_ID) / loss_divisor).sum()
+        q_halt_loss = F.binary_cross_entropy_with_logits(outputs["q_halt_logits"], seq_is_correct.to(outputs["q_halt_logits"].dtype), reduction="sum")
+
+        metrics.update({
+            "lm_loss": lm_loss.detach(),
+            "q_halt_loss": q_halt_loss.detach(),
+        })
+
+        # Q continue (bootstrapping target loss)
+        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()
diff --git a/sudoku-extreme/original-hrm/step_10416 b/sudoku-extreme/original-hrm/step_10416
new file mode 100644
index 0000000000000000000000000000000000000000..950113892b3a1890c4a90d383d967236272ed1cf
--- /dev/null
+++ b/sudoku-extreme/original-hrm/step_10416
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:16e25270e30091f26cbd7e5553031f7b67e77f8f86f807a82c3daa26b76777c3
+size 109124341
diff --git a/sudoku-extreme/original-hrm/step_13020 b/sudoku-extreme/original-hrm/step_13020
new file mode 100644
index 0000000000000000000000000000000000000000..bc58b3a6e32e4d6b1d5a85629d1caaffbb4f6e5d
--- /dev/null
+++ b/sudoku-extreme/original-hrm/step_13020
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:57a9bffe9384bd46fd7caa03a3af8e4a715e877e169f3159f072b213c50ef9af
+size 109124341
diff --git a/sudoku-extreme/original-hrm/step_15624 b/sudoku-extreme/original-hrm/step_15624
new file mode 100644
index 0000000000000000000000000000000000000000..cc77fee76c64eadad3f523b9ed6474d23beafa1c
--- /dev/null
+++ b/sudoku-extreme/original-hrm/step_15624
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ee8c5fbd8880684d24d6b3ebfa3586242fc6a560f9a0eb58c3c8d6f94a7c9503
+size 109124341
diff --git a/sudoku-extreme/original-hrm/step_18228 b/sudoku-extreme/original-hrm/step_18228
new file mode 100644
index 0000000000000000000000000000000000000000..6c557e3396491d531cd87e72388aeabf3bbfd86a
--- /dev/null
+++ b/sudoku-extreme/original-hrm/step_18228
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e2e37f26f20fbb87f9d7198c71f6b736e39b378035793c6dae830761185a22de
+size 109124341
diff --git a/sudoku-extreme/original-hrm/step_20832 b/sudoku-extreme/original-hrm/step_20832
new file mode 100644
index 0000000000000000000000000000000000000000..ad22b5602df51e15fa04f397fc0fc85d99377bc3
--- /dev/null
+++ b/sudoku-extreme/original-hrm/step_20832
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8039d848fca35d2e0d2e674d57cc0e37ec3b2fa4d9aa3fac5e4d9ad4a4cd69de
+size 109124341
diff --git a/sudoku-extreme/original-hrm/step_23436 b/sudoku-extreme/original-hrm/step_23436
new file mode 100644
index 0000000000000000000000000000000000000000..dee66f428b34bf26afbbec489bc666224e150609
--- /dev/null
+++ b/sudoku-extreme/original-hrm/step_23436
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c74f4f32f8ddd2f898be91f89c15494f5216e6fd1938a015b5d364eb2fa415ef
+size 109124341
diff --git a/sudoku-extreme/original-hrm/step_2604 b/sudoku-extreme/original-hrm/step_2604
new file mode 100644
index 0000000000000000000000000000000000000000..f4826efeaf84409444a78438fb420db6274a0678
--- /dev/null
+++ b/sudoku-extreme/original-hrm/step_2604
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ea37d7b836e58fba35df4c055a835920bd5d400052c5c5e04e81b76ecfbccd59
+size 109124296
diff --git a/sudoku-extreme/original-hrm/step_26040 b/sudoku-extreme/original-hrm/step_26040
new file mode 100644
index 0000000000000000000000000000000000000000..487a5c9b072c37f9a52ae044d65a51d0964bc4e8
--- /dev/null
+++ b/sudoku-extreme/original-hrm/step_26040
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:800399ff98efd2f0e638fcc5bc8071823e2e47d0e4625ed0061d9b801d2b9c16
+size 109124341
diff --git a/sudoku-extreme/original-hrm/step_28644 b/sudoku-extreme/original-hrm/step_28644
new file mode 100644
index 0000000000000000000000000000000000000000..fe462b384415d46f845e80008f17c3febda14545
--- /dev/null
+++ b/sudoku-extreme/original-hrm/step_28644
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f89e5d894ad10c064de31542e765a5be68a29c3dbe71f5f631d288ef6e32da3f
+size 109124341
diff --git a/sudoku-extreme/original-hrm/step_31248 b/sudoku-extreme/original-hrm/step_31248
new file mode 100644
index 0000000000000000000000000000000000000000..55dd68c1d9d44e83ffebbfe9cb0c1f7b2db4ff01
--- /dev/null
+++ b/sudoku-extreme/original-hrm/step_31248
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d50339858d376424b1b3a39162f347c16f608f8502db188248d4985a84241f81
+size 109124341
diff --git a/sudoku-extreme/original-hrm/step_33852 b/sudoku-extreme/original-hrm/step_33852
new file mode 100644
index 0000000000000000000000000000000000000000..80d34d0a41e689ccf27e849cc1748f6f3c79a46e
--- /dev/null
+++ b/sudoku-extreme/original-hrm/step_33852
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3153792dacd5bec98744ae866f219b21b1aa5dd417470ab8aef8622ed5225506
+size 109124341
diff --git a/sudoku-extreme/original-hrm/step_36456 b/sudoku-extreme/original-hrm/step_36456
new file mode 100644
index 0000000000000000000000000000000000000000..2ea50602ba2cb0613893679a5e36910ee381088d
--- /dev/null
+++ b/sudoku-extreme/original-hrm/step_36456
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:342368edd345839d6606df7108f78bdc4d25cab4020bc97cb2eb7e40bcccb45f
+size 109124341
diff --git a/sudoku-extreme/original-hrm/step_39060 b/sudoku-extreme/original-hrm/step_39060
new file mode 100644
index 0000000000000000000000000000000000000000..e3190ff2d41372368e2d854e97859c5edf879896
--- /dev/null
+++ b/sudoku-extreme/original-hrm/step_39060
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9eb1da2e8a6b42eaade47ee3f6880c7ae87bf670a7145252d6059ea0cf876820
+size 109124341
diff --git a/sudoku-extreme/original-hrm/step_41664 b/sudoku-extreme/original-hrm/step_41664
new file mode 100644
index 0000000000000000000000000000000000000000..d4cb15c2afbac58961e4f7995d4d18f80ca68b20
--- /dev/null
+++ b/sudoku-extreme/original-hrm/step_41664
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:64cbe88964f442cf2e5b2b0503b24c1d8e03ee6a5a274ee883ab23f08a17955c
+size 109124341
diff --git a/sudoku-extreme/original-hrm/step_44268 b/sudoku-extreme/original-hrm/step_44268
new file mode 100644
index 0000000000000000000000000000000000000000..712020d14f5aaa31355ce7005f17844b52573bdb
--- /dev/null
+++ b/sudoku-extreme/original-hrm/step_44268
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8bd867849981936ddc1d3cd18a36946479c49c5ee08f17c08faa7bd62afcf1bb
+size 109124341
diff --git a/sudoku-extreme/original-hrm/step_46872 b/sudoku-extreme/original-hrm/step_46872
new file mode 100644
index 0000000000000000000000000000000000000000..776d620adcdb8c0e78a675f263e7cfc9bea3a471
--- /dev/null
+++ b/sudoku-extreme/original-hrm/step_46872
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e16c83e84eef00db909d46ba52f060b2650d924b178d3d721a2e7116ce2b9bf9
+size 109124341
diff --git a/sudoku-extreme/original-hrm/step_49476 b/sudoku-extreme/original-hrm/step_49476
new file mode 100644
index 0000000000000000000000000000000000000000..dd48e38c0e8d067ec9f90a13697e92e9bff2989f
--- /dev/null
+++ b/sudoku-extreme/original-hrm/step_49476
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:91f0c729499f820bba7da5af462bbbf4afb19b5dee04b57a345e08e98bdf60a5
+size 109124341
diff --git a/sudoku-extreme/original-hrm/step_5208 b/sudoku-extreme/original-hrm/step_5208
new file mode 100644
index 0000000000000000000000000000000000000000..381ec056d9c875971c1f10d11379d2c3a562f0cb
--- /dev/null
+++ b/sudoku-extreme/original-hrm/step_5208
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fa3ceb388d0aa4c0deadac5e6b110a22b743ba6dcae62b20bb97e7bf95000d1a
+size 109124296
diff --git a/sudoku-extreme/original-hrm/step_52080 b/sudoku-extreme/original-hrm/step_52080
new file mode 100644
index 0000000000000000000000000000000000000000..642d2acb54ab413472da0ebf14e18edef53a2c5b
--- /dev/null
+++ b/sudoku-extreme/original-hrm/step_52080
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5e8703bc491deb2527ac7943beda76cb6041d4e9e3b2f52829cd3d5f2adbca53
+size 109124341
diff --git a/sudoku-extreme/original-hrm/step_7812 b/sudoku-extreme/original-hrm/step_7812
new file mode 100644
index 0000000000000000000000000000000000000000..4a04abb788eb4a5c57ddb1f2c4c4a6d29a6bad4b
--- /dev/null
+++ b/sudoku-extreme/original-hrm/step_7812
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cc40729bbfd3f4f99e971ac9c5ec5f1c5fc3fe75d7b2de7f9044aba8fe4502c6
+size 109124296