diff --git a/.gitattributes b/.gitattributes
index f735ecff52204e8840b0ea9079a47c1c962f564c..dcd9e2a2294a2909a4257b69a1b9a083183e0170 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -253,3 +253,80 @@ sudoku-extreme/trm-attention/step_52080 filter=lfs diff=lfs merge=lfs -text
 sudoku-extreme/trm-attention/step_6510 filter=lfs diff=lfs merge=lfs -text
 sudoku-extreme/trm-attention/step_7812 filter=lfs diff=lfs merge=lfs -text
 sudoku-extreme/trm-attention/step_9114 filter=lfs diff=lfs merge=lfs -text
+maze-hard/original-hrm/step_101556 filter=lfs diff=lfs merge=lfs -text
+maze-hard/original-hrm/step_109368 filter=lfs diff=lfs merge=lfs -text
+maze-hard/original-hrm/step_117180 filter=lfs diff=lfs merge=lfs -text
+maze-hard/original-hrm/step_124992 filter=lfs diff=lfs merge=lfs -text
+maze-hard/original-hrm/step_132804 filter=lfs diff=lfs merge=lfs -text
+maze-hard/original-hrm/step_140616 filter=lfs diff=lfs merge=lfs -text
+maze-hard/original-hrm/step_148428 filter=lfs diff=lfs merge=lfs -text
+maze-hard/original-hrm/step_15624 filter=lfs diff=lfs merge=lfs -text
+maze-hard/original-hrm/step_156240 filter=lfs diff=lfs merge=lfs -text
+maze-hard/original-hrm/step_23436 filter=lfs diff=lfs merge=lfs -text
+maze-hard/original-hrm/step_31248 filter=lfs diff=lfs merge=lfs -text
+maze-hard/original-hrm/step_39060 filter=lfs diff=lfs merge=lfs -text
+maze-hard/original-hrm/step_46872 filter=lfs diff=lfs merge=lfs -text
+maze-hard/original-hrm/step_54684 filter=lfs diff=lfs merge=lfs -text
+maze-hard/original-hrm/step_62496 filter=lfs diff=lfs merge=lfs -text
+maze-hard/original-hrm/step_70308 filter=lfs diff=lfs merge=lfs -text
+maze-hard/original-hrm/step_7812 filter=lfs diff=lfs merge=lfs -text
+maze-hard/original-hrm/step_78120 filter=lfs diff=lfs merge=lfs -text
+maze-hard/original-hrm/step_85932 filter=lfs diff=lfs merge=lfs -text
+maze-hard/original-hrm/step_93744 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-large/step_101556 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-large/step_109368 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-large/step_117180 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-large/step_124992 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-large/step_132804 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-large/step_140616 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-large/step_148428 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-large/step_15624 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-large/step_156240 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-large/step_23436 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-large/step_31248 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-large/step_39060 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-large/step_46872 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-large/step_54684 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-large/step_62496 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-large/step_70308 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-large/step_7812 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-large/step_78120 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-large/step_85932 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-large/step_93744 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_101556 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_10416 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_109368 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_11718 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_117180 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_124992 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_1302 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_13020 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_132804 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_140616 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_14322 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_148428 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_15624 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_156240 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_16926 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_18228 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_19530 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_20832 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_22134 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_23436 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_24738 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_2604 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_26040 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_31248 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_3906 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_39060 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_46872 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_5208 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_54684 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_62496 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_6510 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_70308 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_7812 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_78120 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_85932 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_9114 filter=lfs diff=lfs merge=lfs -text
+maze-hard/shrek-tiny/step_93744 filter=lfs diff=lfs merge=lfs -text
diff --git a/maze-hard/original-hrm/all_config.yaml b/maze-hard/original-hrm/all_config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..662d832747b33b3d9abfc5bd291213de79690ccc
--- /dev/null
+++ b/maze-hard/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/HRM_Maze_Comparison/HierarchicalReasoningModel_ACTV1
+  stalwart-oryx
+data_path: ../../../dataset/data/maze-30x30-hard-1k
+epochs: 20000
+eval_interval: 1000
+eval_save_outputs: []
+global_batch_size: 128
+lr: 0.0001
+lr_min_ratio: 1.0
+lr_warmup_steps: 2000
+project_name: HRM_Maze_Comparison
+puzzle_emb_lr: 0.0001
+puzzle_emb_weight_decay: 1.0
+run_name: HierarchicalReasoningModel_ACTV1 stalwart-oryx
+seed: 0
+weight_decay: 1.0
diff --git a/maze-hard/original-hrm/hrm_act_v1.py b/maze-hard/original-hrm/hrm_act_v1.py
new file mode 100644
index 0000000000000000000000000000000000000000..e91c7d1a047b9f9735fa8253daa4b8effa23f011
--- /dev/null
+++ b/maze-hard/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/maze-hard/original-hrm/losses.py b/maze-hard/original-hrm/losses.py
new file mode 100644
index 0000000000000000000000000000000000000000..b3118e728365cc2824646a156a35d79f6dcca9c1
--- /dev/null
+++ b/maze-hard/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/maze-hard/original-hrm/step_101556 b/maze-hard/original-hrm/step_101556
new file mode 100644
index 0000000000000000000000000000000000000000..cd0e389176e1fa33636651539c5afc5a5572e754
--- /dev/null
+++ b/maze-hard/original-hrm/step_101556
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ef9c68a6903d5dedd80c0784ae8312a0824ad2ee9a10521809b7b012dc00d2ae
+size 109103906
diff --git a/maze-hard/original-hrm/step_109368 b/maze-hard/original-hrm/step_109368
new file mode 100644
index 0000000000000000000000000000000000000000..44c87df1e607b668c53b9d2118265bfaf87a9c03
--- /dev/null
+++ b/maze-hard/original-hrm/step_109368
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:10dd5a1e8682f17fba4f09863376bc4452040b5b1dfe7dd9f403870e58ffa6ee
+size 109103906
diff --git a/maze-hard/original-hrm/step_117180 b/maze-hard/original-hrm/step_117180
new file mode 100644
index 0000000000000000000000000000000000000000..9feb34fb9bd3255797596f724ca0a593275b97a2
--- /dev/null
+++ b/maze-hard/original-hrm/step_117180
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:370df57b9bcd3a8430206a1441a264dc1a76403dbebdbc70e84d958b19611ee6
+size 109103906
diff --git a/maze-hard/original-hrm/step_124992 b/maze-hard/original-hrm/step_124992
new file mode 100644
index 0000000000000000000000000000000000000000..7cd42b155b53ba8df738d176504e01fb3a811ce9
--- /dev/null
+++ b/maze-hard/original-hrm/step_124992
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:50258c4fd7e2343a52bfcadb221805721e7e91a71cd24373dbdb7e9ab5389dc3
+size 109103906
diff --git a/maze-hard/original-hrm/step_132804 b/maze-hard/original-hrm/step_132804
new file mode 100644
index 0000000000000000000000000000000000000000..af97d5349da289ee86f2a9dcbf6ed8fec5cc040c
--- /dev/null
+++ b/maze-hard/original-hrm/step_132804
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3449afe9f7d8e2068d8cecb92c711f0cefa15620433d0d8ccfd82bbb4bf2b042
+size 109103906
diff --git a/maze-hard/original-hrm/step_140616 b/maze-hard/original-hrm/step_140616
new file mode 100644
index 0000000000000000000000000000000000000000..c29dec714a0163c98276846e7cd1e7f9572d78d9
--- /dev/null
+++ b/maze-hard/original-hrm/step_140616
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0db23d79c94ef4553e76f191127c8a484a07c3bb58ab2134e42dd8fb8493c511
+size 109103906
diff --git a/maze-hard/original-hrm/step_148428 b/maze-hard/original-hrm/step_148428
new file mode 100644
index 0000000000000000000000000000000000000000..44699cb3367b7ed5db9cd7ea1636213ca7c0ca3f
--- /dev/null
+++ b/maze-hard/original-hrm/step_148428
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:950498283faa9832b2fb0da91b4975cd04b75cba91320d01b4c55922c72974c7
+size 109103906
diff --git a/maze-hard/original-hrm/step_15624 b/maze-hard/original-hrm/step_15624
new file mode 100644
index 0000000000000000000000000000000000000000..f8741f554811c76339afb3ec1973a92db1440cf9
--- /dev/null
+++ b/maze-hard/original-hrm/step_15624
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:24433ec581392dc0677b1a0b0b3f50af16edecf1b5b38e1e6a06b82a3bdcea59
+size 109103861
diff --git a/maze-hard/original-hrm/step_156240 b/maze-hard/original-hrm/step_156240
new file mode 100644
index 0000000000000000000000000000000000000000..00c4160bc799866b0f6138e88db3e56db6dfc9ab
--- /dev/null
+++ b/maze-hard/original-hrm/step_156240
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:148d8bd84ed22d4db1eddef76dab227d6e5fe6955de7cb8810248ccc98ebeac3
+size 109103906
diff --git a/maze-hard/original-hrm/step_23436 b/maze-hard/original-hrm/step_23436
new file mode 100644
index 0000000000000000000000000000000000000000..c0cae14c36d6bf5b894a26d486c07bd8986969b3
--- /dev/null
+++ b/maze-hard/original-hrm/step_23436
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e7dbe385f6e9c11c3141391471a4b4cd79dadaff4738b8ea0f8fc47dc9dd17cb
+size 109103861
diff --git a/maze-hard/original-hrm/step_31248 b/maze-hard/original-hrm/step_31248
new file mode 100644
index 0000000000000000000000000000000000000000..d0d718c7263b4f337be7abc6f5518be6f6bdb678
--- /dev/null
+++ b/maze-hard/original-hrm/step_31248
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:232903a94774281767b2bcbcd4579d8a32c822816e8646ca2e1ee67110997e5f
+size 109103861
diff --git a/maze-hard/original-hrm/step_39060 b/maze-hard/original-hrm/step_39060
new file mode 100644
index 0000000000000000000000000000000000000000..48fe83dab9240683b279f78efcf75ab2412eda57
--- /dev/null
+++ b/maze-hard/original-hrm/step_39060
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c1a956113527727345c0d636a9cb05c007affd0b3f6da752ec542bc74412dd79
+size 109103861
diff --git a/maze-hard/original-hrm/step_46872 b/maze-hard/original-hrm/step_46872
new file mode 100644
index 0000000000000000000000000000000000000000..f30d239f786dc431d0a67eb04551393c323fff55
--- /dev/null
+++ b/maze-hard/original-hrm/step_46872
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f4d7810b5fc2e97138f9b2320a81c0af3f94c0e876ea505526beb23d4f0a3c0c
+size 109103861
diff --git a/maze-hard/original-hrm/step_54684 b/maze-hard/original-hrm/step_54684
new file mode 100644
index 0000000000000000000000000000000000000000..7b77d4643a96765c2a0598eb9da1f61237112e78
--- /dev/null
+++ b/maze-hard/original-hrm/step_54684
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e57bf7870ceb4f2e19179b29cd041761e6fb891fb80894a2ebcea88b3d9b8bb3
+size 109103861
diff --git a/maze-hard/original-hrm/step_62496 b/maze-hard/original-hrm/step_62496
new file mode 100644
index 0000000000000000000000000000000000000000..8c48a3cd8e2b718904fab66f30940d8b8e9f582e
--- /dev/null
+++ b/maze-hard/original-hrm/step_62496
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2bde8cbff83de2a34fe71bb99d21d1f18c857115acb9e2010026efad8a95918d
+size 109103861
diff --git a/maze-hard/original-hrm/step_70308 b/maze-hard/original-hrm/step_70308
new file mode 100644
index 0000000000000000000000000000000000000000..20e0ec7a5c2fb17de0399b5b9b85bdef7d59bc20
--- /dev/null
+++ b/maze-hard/original-hrm/step_70308
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:da0b19a21a631324c2e00a1297f831f1b2f58c71068704fb4419b2a78ad355c0
+size 109103861
diff --git a/maze-hard/original-hrm/step_7812 b/maze-hard/original-hrm/step_7812
new file mode 100644
index 0000000000000000000000000000000000000000..cf9d155e2865387c35eacf676bd84624aefbb640
--- /dev/null
+++ b/maze-hard/original-hrm/step_7812
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2e21f2f1c004ea64a081fee1763fa5e09631de1365298f59df92d345c572a576
+size 109103816
diff --git a/maze-hard/original-hrm/step_78120 b/maze-hard/original-hrm/step_78120
new file mode 100644
index 0000000000000000000000000000000000000000..ff27da3b48c157249854dac7779b49f7e685c70a
--- /dev/null
+++ b/maze-hard/original-hrm/step_78120
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb34d5ffea78ced7f3985168ec20a4de29ef48f88d5bfe6767b77a009ad2fe76
+size 109103861
diff --git a/maze-hard/original-hrm/step_85932 b/maze-hard/original-hrm/step_85932
new file mode 100644
index 0000000000000000000000000000000000000000..95b1c421ec623e6d21f5fe0875b280d62ae56127
--- /dev/null
+++ b/maze-hard/original-hrm/step_85932
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:50fa60917de1909525681da60c4403a83f6a28d45d5473a9df94579a6d449bdf
+size 109103861
diff --git a/maze-hard/original-hrm/step_93744 b/maze-hard/original-hrm/step_93744
new file mode 100644
index 0000000000000000000000000000000000000000..e5ec8549b0f4c265e2079e868457e75a2213de4b
--- /dev/null
+++ b/maze-hard/original-hrm/step_93744
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:58c046141e6c7b890982e0b78dfdecf8ab788007ea93c9472c736fff61f89730
+size 109103861
diff --git a/maze-hard/shrek-large/all_config.yaml b/maze-hard/shrek-large/all_config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..b84b11bd05cfedc23d8541ae7db02ee383e5a8ac
--- /dev/null
+++ b/maze-hard/shrek-large/all_config.yaml
@@ -0,0 +1,36 @@
+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/HRM_Maze_Comparison/SHREK_Large_Maze
+data_path: ../../dataset/data/maze-30x30-hard-1k
+ema: true
+ema_rate: 0.999
+epochs: 20000
+eval_interval: 1000
+eval_save_outputs: []
+global_batch_size: 128
+lr: 0.0001
+lr_min_ratio: 1.0
+lr_warmup_steps: 2000
+project_name: HRM_Maze_Comparison
+puzzle_emb_lr: 0.0001
+puzzle_emb_weight_decay: 1.0
+run_name: SHREK_Large_Maze
+seed: 0
+weight_decay: 1.0
diff --git a/maze-hard/shrek-large/hrm_act_v1.py b/maze-hard/shrek-large/hrm_act_v1.py
new file mode 100644
index 0000000000000000000000000000000000000000..07f7bfb5b13cd314a68d8e586ddac80fe9182006
--- /dev/null
+++ b/maze-hard/shrek-large/hrm_act_v1.py
@@ -0,0 +1,428 @@
+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
+from models.hrm.error_singals import get_error_signal  # SHREK: error signal module
+
+
+@dataclass
+class HierarchicalReasoningModel_ACTV1InnerCarry:
+    z_H: torch.Tensor
+    z_L: torch.Tensor
+    # SHREK: prev_pred stores last step's argmax predictions for flip rate computation.
+    # zeros = fresh start (first step after init or reset gives flip_rate ≈ 1.0)
+    prev_pred: torch.Tensor        # (B, seq_len) int32
+    # SHREK: Q-values cached in carry — no longer used for Q-targets (see Bug 4),
+    # but kept because removing them causes torch.compile regression.
+    prev_q_halt: torch.Tensor      # (B,) float32
+    prev_q_continue: torch.Tensor  # (B,) float32
+
+
+@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
+
+    # SHREK: error injection warmup — ramps alpha from 0 to alpha_max over warmup steps.
+    # Prevents small models from collapsing before the error estimator is accurate.
+    alpha_max: float = 0.01
+    alpha_warmup_steps: int = 5000
+
+    # SHREK ablation flags — disable components from training scripts
+    enable_error_injection: bool = True    # Error-Conditioned Injection (flip rate + learned estimator)
+    enable_stagnation_delta: bool = True   # Stagnation-Aware Q-head (delta fed to Q-head)
+
+    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)
+        # SHREK: Q-head input size depends on stagnation delta flag
+        q_head_input_size = self.config.hidden_size + (1 if self.config.enable_stagnation_delta else 0)
+        self.q_head       = CastedLinear(q_head_input_size, 2, bias=True)
+
+        # SHREK: error_encoder maps the scalar error score -> hidden_size vector for injection into z_H
+        # alpha follows a linear warmup schedule (0 → alpha_max over warmup steps).
+        # This lets the error estimator train before its signal affects z_H.
+        self.error_encoder  = nn.Linear(1, self.config.hidden_size)
+        # SHREK: step counter for alpha warmup (not a learned parameter)
+        self.register_buffer('_alpha_step', torch.tensor(0, dtype=torch.long))
+        # SHREK: error_estimator reads z_H and predicts how wrong the model is.
+        # trained via auxiliary loss in pretrain.py using the real lm_loss as target.
+        # catches "stuck but wrong" — a model confidently on the wrong answer.
+        # flip rate catches oscillation; estimator catches confident-but-wrong.
+        self.error_estimator = nn.Linear(self.config.hidden_size, 1)
+
+        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),
+            # SHREK: zeros = no previous prediction — first step gives flip_rate ≈ 1.0
+            # device=H_init.device ensures prev_pred is on CUDA, matching z_H and logits
+            prev_pred=torch.zeros(batch_size, self.config.seq_len, dtype=torch.int32, device=self.H_init.device),
+            prev_q_halt=torch.full((batch_size,), -5.0, device=self.H_init.device),
+            prev_q_continue=torch.full((batch_size,), -5.0, device=self.H_init.device),
+        )
+        
+    def reset_carry(self, reset_flag: torch.Tensor, carry: HierarchicalReasoningModel_ACTV1InnerCarry):
+        # SHREK: zero out prev_pred for reset sequences so they start fresh.
+        # a reset sequence is one that just halted — it will solve a new puzzle next.
+        # zeroing prev_pred means first step gives flip_rate ≈ 1.0 (maximum uncertainty).
+        new_prev_pred = torch.where(
+            reset_flag.view(-1, 1),
+            torch.zeros_like(carry.prev_pred),
+            carry.prev_pred
+        )
+        new_prev_q_halt     = torch.where(reset_flag, torch.full_like(carry.prev_q_halt,     -5.0), carry.prev_q_halt)
+        new_prev_q_continue = torch.where(reset_flag, torch.full_like(carry.prev_q_continue, -5.0), carry.prev_q_continue)
+
+        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),
+            prev_pred=new_prev_pred,
+            prev_q_halt=new_prev_q_halt,
+            prev_q_continue=new_prev_q_continue,
+        )
+
+
+    # SHREK: removed task_type parameter — error signal is now universal (no task rules needed)
+    def forward(self, carry: HierarchicalReasoningModel_ACTV1InnerCarry, batch: Dict[str, torch.Tensor], require_trace=False):
+    #  -> Tuple[HierarchicalReasoningModel_ACTV1InnerCarry, torch.Tensor, Tuple[torch.Tensor, 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 — decode z_H into token predictions
+        output = self.lm_head(z_H)[:, self.puzzle_emb_len:]   # (B, seq_len, vocab_size)
+
+        # SHREK Component 1: Combined Error Signal
+        # Signal A — flip rate: what fraction of tokens changed from last step?
+        #   catches oscillation (model changing its mind between wrong options)
+        flip_err, current_pred = get_error_signal(output, carry.prev_pred)     # (B,), (B, seq_len)
+
+        # Signal B — learned estimator: reads z_H and predicts how wrong the model is.
+        #   catches stuck-but-wrong (model confidently on wrong answer without oscillating)
+        #   average over content positions (skip puzzle embedding prefix positions)
+        #   CRITICAL: detach z_H_mean so aux_loss only trains error_estimator weights,
+        #   NOT the main reasoning layers. Without detach, aux_loss sends parasitic
+        #   gradients through z_H → H_level/L_level that conflict with lm_loss.
+        z_H_mean    = z_H[:, self.puzzle_emb_len:].mean(dim=1).detach()        # (B, hidden_size)
+        learned_err = torch.sigmoid(self.error_estimator(z_H_mean.float()))    # (B, 1)
+        learned_err = learned_err.squeeze(-1)                                  # (B,)
+
+        # SHREK: combined error = 50/50 blend of both signals
+        # flip_err works immediately from step 1 (no learning needed)
+        # learned_err becomes accurate over training and takes over as the stronger signal
+        error = 0.5 * flip_err + 0.5 * learned_err                            # (B,)
+
+        # SHREK Component 2: Stagnation-Aware Q-head
+        # Ablation: can be disabled via config.enable_stagnation_delta=False
+        if self.config.enable_stagnation_delta:
+            # Compute how much z_H changed during this reasoning step.
+            # CRITICAL: detach so Q-head gradients don't flow through all positions of z_H.
+            # Without detach, Q-head loss would route gradients through all 82 positions,
+            # drowning out lm_loss. With detach, Q-head gradients only flow through position 0.
+            z_H_f = z_H.detach().float()
+            z_H_start = carry.z_H.detach().float()
+            delta = torch.norm(z_H_f - z_H_start, dim=(1, 2)) / (torch.norm(z_H_start, dim=(1, 2)) + 1e-6)  # (B,)
+            # Q-head reads CLS token + stagnation delta
+            q_input = torch.cat([z_H[:, 0].to(torch.float32), delta.unsqueeze(-1)], dim=-1)  # (B, hidden_size+1)
+        else:
+            # Fallback: same as original HRM — Q-head reads CLS token only
+            q_input = z_H[:, 0].to(torch.float32)  # (B, hidden_size)
+        q_logits = self.q_head(q_input).to(torch.float32)  # (B, 2)
+
+        # SHREK: inject combined error into z_H (AFTER Q-head, only affects carry for next step)
+        # Ablation: can be disabled via config.enable_error_injection=False
+        if self.config.enable_error_injection:
+            # error_encoder maps scalar -> hidden_size vector
+            # alpha follows linear warmup: 0 → alpha_max over warmup steps
+            # scaled by 1/sqrt(hidden_size) so injection is proportional to model size
+            error_emb = self.error_encoder(error.unsqueeze(-1))                    # (B, hidden_size)
+            # SHREK: compute alpha from warmup schedule (not learned)
+            # During warmup, alpha ramps linearly from 0 to alpha_max.
+            # After warmup, alpha stays at alpha_max.
+            # Uses torch.clamp instead of Python min() to stay compatible with torch.compile.
+            with torch.no_grad():
+                if self.training:
+                    self._alpha_step += 1
+                alpha = self.config.alpha_max * torch.clamp(self._alpha_step / self.config.alpha_warmup_steps, max=1.0)
+            scale = math.sqrt(self.config.hidden_size)
+            z_H = z_H + (alpha * error_emb.unsqueeze(1) / scale).to(z_H.dtype)   # (B, seq_len, hidden_size)
+
+        # New carry: store error-injected z_H so next ACT step starts from it
+        new_carry = HierarchicalReasoningModel_ACTV1InnerCarry(
+            z_H=z_H.detach(),
+            z_L=z_L.detach(),
+            # SHREK: store current predictions — next step compares against these for flip rate
+            prev_pred=current_pred.detach(),
+            # SHREK: Q-values written to carry for torch.compile compatibility (not read for Q-targets)
+            prev_q_halt=q_logits[..., 0].detach(),
+            prev_q_continue=q_logits[..., 1].detach(),
+        )
+
+        # SHREK: also return learned_err so pretrain.py can compute auxiliary loss
+        # auxiliary loss trains the estimator: predicted error should match real lm_loss
+        if require_trace:
+            return z_H_trace, new_carry, output, (q_logits[..., 0], q_logits[..., 1]), learned_err
+        else:
+            return new_carry, output, (q_logits[..., 0], q_logits[..., 1]), learned_err
+
+
+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()}
+        )
+        
+    # SHREK: removed task_type parameter — error signal is universal, no task rules needed
+    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()}
+
+        # SHREK: run inner forward — unpack learned_err for auxiliary loss in pretrain.py
+        if require_trace:
+            z_H_trace, new_inner_carry, logits, (q_halt_logits, q_continue_logits), learned_err = self.inner(new_inner_carry, new_current_data, require_trace=require_trace)
+        else:
+            new_inner_carry, logits, (q_halt_logits, q_continue_logits), learned_err = self.inner(new_inner_carry, new_current_data)
+
+        outputs = {
+            "logits": logits,
+            "q_halt_logits": q_halt_logits,
+            "q_continue_logits": q_continue_logits,
+            "learned_err": learned_err,   # SHREK: for auxiliary loss in pretrain.py
+        }
+        
+        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)
+
+                # SHREK: Q-target via double forward pass (same as original HRM).
+                # Run inner() a second time to get NEXT step's Q-values (step T+1).
+                # Save/restore _alpha_step so the second call doesn't double-count
+                # the warmup — without this, alpha reaches full strength at step ~2500
+                # instead of 5000, destabilizing early training.
+                # NOTE: prev_q fields in carry are not read here — they must stay
+                # in the dataclass for torch.compile compatibility.
+                saved_alpha_step = self.inner._alpha_step.clone()
+                next_q_halt_logits, next_q_continue_logits = self.inner(new_inner_carry, new_current_data)[-2]
+                self.inner._alpha_step.copy_(saved_alpha_step)
+
+                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/maze-hard/shrek-large/losses.py b/maze-hard/shrek-large/losses.py
new file mode 100644
index 0000000000000000000000000000000000000000..784535de2c64479d78806a5a3d48f24b96ce16e2
--- /dev/null
+++ b/maze-hard/shrek-large/losses.py
@@ -0,0 +1,130 @@
+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()
+
+        # SHREK: auxiliary loss — trains error_estimator to predict per-sample lm_loss.
+        # Computed here (not in pretrain.py) so learned_err still has gradients.
+        # per_sample_lm_loss is detached so only error_estimator weights get updated by aux_loss.
+        aux_loss = 0
+        if "learned_err" in outputs:
+            per_sample_lm_loss = (ds_mask * loss / loss_divisor).sum(dim=1).detach()  # (B,)
+            # SHREK: min-max normalize target to 0-1 so it matches sigmoid output range.
+            # When all losses are equal (early training), target = 0 (no distinguishable error).
+            # When losses vary, target spreads across 0-1 giving the estimator a useful signal.
+            lm_min = per_sample_lm_loss.min()
+            lm_max = per_sample_lm_loss.max()
+            target = (per_sample_lm_loss - lm_min) / (lm_max - lm_min + 1e-6)       # (B,) in [0, 1]
+            aux_loss = F.mse_loss(outputs["learned_err"], target, reduction="sum")
+            metrics["aux_loss"] = aux_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) + 0.1 * aux_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) + 0.1 * aux_loss, metrics, detached_outputs, new_carry.halted.all(), new_carry.halted & act_halt
diff --git a/maze-hard/shrek-large/step_101556 b/maze-hard/shrek-large/step_101556
new file mode 100644
index 0000000000000000000000000000000000000000..1054655d70c937bd9426ddcd48dca0730abc5595
--- /dev/null
+++ b/maze-hard/shrek-large/step_101556
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08ccbff96ecfe43f613e4ebe70bbe16dd92f09439edd28808ca3fc4519b3e278
+size 109111847
diff --git a/maze-hard/shrek-large/step_109368 b/maze-hard/shrek-large/step_109368
new file mode 100644
index 0000000000000000000000000000000000000000..63fb32a08a045f77951bd6899efd82a3dbfe58bf
--- /dev/null
+++ b/maze-hard/shrek-large/step_109368
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3d7d6fe90c6e1fad34b7dd334c36b5342b4412541239f7cf0199954e381224a4
+size 109111847
diff --git a/maze-hard/shrek-large/step_117180 b/maze-hard/shrek-large/step_117180
new file mode 100644
index 0000000000000000000000000000000000000000..9dd543163e343684abad5b635124dabd1bdc523f
--- /dev/null
+++ b/maze-hard/shrek-large/step_117180
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:71b3c0f56ff3fd5afdc05dfb8bde57b43a193b182825b2c246e498b6911c2379
+size 109111847
diff --git a/maze-hard/shrek-large/step_124992 b/maze-hard/shrek-large/step_124992
new file mode 100644
index 0000000000000000000000000000000000000000..69b1c23943e2102c2407d57e075ecb688583c42a
--- /dev/null
+++ b/maze-hard/shrek-large/step_124992
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:326936b97c74e3c1bbd975a8bfb8a076db7c1bdfa96cc91f5181ee893625c0be
+size 109111847
diff --git a/maze-hard/shrek-large/step_132804 b/maze-hard/shrek-large/step_132804
new file mode 100644
index 0000000000000000000000000000000000000000..52b01831477deccd17e0668520c3cb3962543001
--- /dev/null
+++ b/maze-hard/shrek-large/step_132804
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:27f1f125b33ab8a44b1d1ef6d5515dc3c0bef0c88df525a231cb1d6752fe5a06
+size 109111847
diff --git a/maze-hard/shrek-large/step_140616 b/maze-hard/shrek-large/step_140616
new file mode 100644
index 0000000000000000000000000000000000000000..611638c82152d2c91baf67bc31b8b9bd7f47a049
--- /dev/null
+++ b/maze-hard/shrek-large/step_140616
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f247a5c5db6eefa268be10afdb517259d51d37419a1566985f7de34adff92b4f
+size 109111847
diff --git a/maze-hard/shrek-large/step_148428 b/maze-hard/shrek-large/step_148428
new file mode 100644
index 0000000000000000000000000000000000000000..ac0f5f43be0f3bf55238d0b0ba6c666fa60ff445
--- /dev/null
+++ b/maze-hard/shrek-large/step_148428
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b37762f33727c3860def9bb511976b18ddaa1e200bef6a06d022375b07e24303
+size 109111847
diff --git a/maze-hard/shrek-large/step_15624 b/maze-hard/shrek-large/step_15624
new file mode 100644
index 0000000000000000000000000000000000000000..8885eab0c2189c1c43f97543721eb7063167756b
--- /dev/null
+++ b/maze-hard/shrek-large/step_15624
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:45e4add6c2dfc906588a48837c10f37563997a8f708e72958014525f96a129d9
+size 109111797
diff --git a/maze-hard/shrek-large/step_156240 b/maze-hard/shrek-large/step_156240
new file mode 100644
index 0000000000000000000000000000000000000000..b45afe76e86374fc4bfa034a2c72031c8e860537
--- /dev/null
+++ b/maze-hard/shrek-large/step_156240
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7cedd60dfbddddf69796dfdeb983ed41446df93c0aedc7bf1aebc9e43425c2a1
+size 109111847
diff --git a/maze-hard/shrek-large/step_23436 b/maze-hard/shrek-large/step_23436
new file mode 100644
index 0000000000000000000000000000000000000000..093465a0ff99a828974ce6ce1b04af707200e7fa
--- /dev/null
+++ b/maze-hard/shrek-large/step_23436
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bf9c10bd299cf10ba281857c92428a387975d71fc4a1d8f39b801cb88d5cafe1
+size 109111797
diff --git a/maze-hard/shrek-large/step_31248 b/maze-hard/shrek-large/step_31248
new file mode 100644
index 0000000000000000000000000000000000000000..6a2ff2ccb843fef0d5a5a583f9cace5655d4fbb6
--- /dev/null
+++ b/maze-hard/shrek-large/step_31248
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b0c59fa16073c0302aa19f7eff997b227f794f72d7190e40611290280792451d
+size 109111797
diff --git a/maze-hard/shrek-large/step_39060 b/maze-hard/shrek-large/step_39060
new file mode 100644
index 0000000000000000000000000000000000000000..fc64ec5abaf40df66aabc7642c83bfa91dff20de
--- /dev/null
+++ b/maze-hard/shrek-large/step_39060
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0d08ad4309fc6595ad92fcd74a4a5da9007d3c836e088739d4a0053423b06537
+size 109111797
diff --git a/maze-hard/shrek-large/step_46872 b/maze-hard/shrek-large/step_46872
new file mode 100644
index 0000000000000000000000000000000000000000..f2e3ad1930ffece53eead0defa4d2cdc8c190a66
--- /dev/null
+++ b/maze-hard/shrek-large/step_46872
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b3451b89f913e209a86a75ca06e1314a8fb0fa5dd7da3809089c8f57350b07f5
+size 109111797
diff --git a/maze-hard/shrek-large/step_54684 b/maze-hard/shrek-large/step_54684
new file mode 100644
index 0000000000000000000000000000000000000000..1fa411928332be6b4a8ac7cfed977dc901be756d
--- /dev/null
+++ b/maze-hard/shrek-large/step_54684
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9a69447080712933e9195514ed7b96c9bc3e6a25d2989cda4fa1540ac1e50aaf
+size 109111797
diff --git a/maze-hard/shrek-large/step_62496 b/maze-hard/shrek-large/step_62496
new file mode 100644
index 0000000000000000000000000000000000000000..64c08aabb3bee91076fc07049257913af6b47859
--- /dev/null
+++ b/maze-hard/shrek-large/step_62496
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7043f51eae15d7c889d8f0d2cfbf1df57a448a1c78397439f7adf9a7def5d14d
+size 109111797
diff --git a/maze-hard/shrek-large/step_70308 b/maze-hard/shrek-large/step_70308
new file mode 100644
index 0000000000000000000000000000000000000000..e64de61163c6cfb1f973104d95d968c1864702f4
--- /dev/null
+++ b/maze-hard/shrek-large/step_70308
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c492ec94fc95b45ced1f7e3e2fefb96fb200f0d7d67046a74b6716115b476497
+size 109111797
diff --git a/maze-hard/shrek-large/step_7812 b/maze-hard/shrek-large/step_7812
new file mode 100644
index 0000000000000000000000000000000000000000..c73c4f3c2c9a3bf23c7b4c2d3d3047dd7c26dde7
--- /dev/null
+++ b/maze-hard/shrek-large/step_7812
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e51f3e00e1a30d22a377ac522120f600b5d8bd0e7c894ce1bbdabc26cae6fe37
+size 109111747
diff --git a/maze-hard/shrek-large/step_78120 b/maze-hard/shrek-large/step_78120
new file mode 100644
index 0000000000000000000000000000000000000000..d697474dbf136062ac31cfa6a9fa866fc34d11f8
--- /dev/null
+++ b/maze-hard/shrek-large/step_78120
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:26835cb1b232c8b52e5b2d539231bb6c13a449fafcd68fe5164854acd33319b5
+size 109111797
diff --git a/maze-hard/shrek-large/step_85932 b/maze-hard/shrek-large/step_85932
new file mode 100644
index 0000000000000000000000000000000000000000..3447c50b5ea74e39d303ad759d0f2e01f6697ab7
--- /dev/null
+++ b/maze-hard/shrek-large/step_85932
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b2dcbf7dd090524a644c088fde7547329017554e728d6fc61d50b4029d7544ec
+size 109111797
diff --git a/maze-hard/shrek-large/step_93744 b/maze-hard/shrek-large/step_93744
new file mode 100644
index 0000000000000000000000000000000000000000..d55afd12050edf2348d6b7b11c536d81bebcf66e
--- /dev/null
+++ b/maze-hard/shrek-large/step_93744
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8382c66d053d0d376a25d8dec4b882a746ea83af9bb4319c6d0bfd4d1506f239
+size 109111797
diff --git a/maze-hard/shrek-tiny/all_config.yaml b/maze-hard/shrek-tiny/all_config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..9e12fbf900f03544de1d4c19664f2ce01676132b
--- /dev/null
+++ b/maze-hard/shrek-tiny/all_config.yaml
@@ -0,0 +1,36 @@
+arch:
+  H_cycles: 2
+  H_layers: 2
+  L_cycles: 2
+  L_layers: 2
+  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/HRM_Maze_Comparison/SHREK_Tiny_Maze
+data_path: ../../dataset/data/maze-30x30-hard-1k
+ema: true
+ema_rate: 0.999
+epochs: 20000
+eval_interval: 1000
+eval_save_outputs: []
+global_batch_size: 128
+lr: 0.0001
+lr_min_ratio: 1.0
+lr_warmup_steps: 2000
+project_name: HRM_Maze_Comparison
+puzzle_emb_lr: 0.0001
+puzzle_emb_weight_decay: 1.0
+run_name: SHREK_Tiny_Maze
+seed: 0
+weight_decay: 1.0
diff --git a/maze-hard/shrek-tiny/hrm_act_v1.py b/maze-hard/shrek-tiny/hrm_act_v1.py
new file mode 100644
index 0000000000000000000000000000000000000000..07f7bfb5b13cd314a68d8e586ddac80fe9182006
--- /dev/null
+++ b/maze-hard/shrek-tiny/hrm_act_v1.py
@@ -0,0 +1,428 @@
+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
+from models.hrm.error_singals import get_error_signal  # SHREK: error signal module
+
+
+@dataclass
+class HierarchicalReasoningModel_ACTV1InnerCarry:
+    z_H: torch.Tensor
+    z_L: torch.Tensor
+    # SHREK: prev_pred stores last step's argmax predictions for flip rate computation.
+    # zeros = fresh start (first step after init or reset gives flip_rate ≈ 1.0)
+    prev_pred: torch.Tensor        # (B, seq_len) int32
+    # SHREK: Q-values cached in carry — no longer used for Q-targets (see Bug 4),
+    # but kept because removing them causes torch.compile regression.
+    prev_q_halt: torch.Tensor      # (B,) float32
+    prev_q_continue: torch.Tensor  # (B,) float32
+
+
+@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
+
+    # SHREK: error injection warmup — ramps alpha from 0 to alpha_max over warmup steps.
+    # Prevents small models from collapsing before the error estimator is accurate.
+    alpha_max: float = 0.01
+    alpha_warmup_steps: int = 5000
+
+    # SHREK ablation flags — disable components from training scripts
+    enable_error_injection: bool = True    # Error-Conditioned Injection (flip rate + learned estimator)
+    enable_stagnation_delta: bool = True   # Stagnation-Aware Q-head (delta fed to Q-head)
+
+    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)
+        # SHREK: Q-head input size depends on stagnation delta flag
+        q_head_input_size = self.config.hidden_size + (1 if self.config.enable_stagnation_delta else 0)
+        self.q_head       = CastedLinear(q_head_input_size, 2, bias=True)
+
+        # SHREK: error_encoder maps the scalar error score -> hidden_size vector for injection into z_H
+        # alpha follows a linear warmup schedule (0 → alpha_max over warmup steps).
+        # This lets the error estimator train before its signal affects z_H.
+        self.error_encoder  = nn.Linear(1, self.config.hidden_size)
+        # SHREK: step counter for alpha warmup (not a learned parameter)
+        self.register_buffer('_alpha_step', torch.tensor(0, dtype=torch.long))
+        # SHREK: error_estimator reads z_H and predicts how wrong the model is.
+        # trained via auxiliary loss in pretrain.py using the real lm_loss as target.
+        # catches "stuck but wrong" — a model confidently on the wrong answer.
+        # flip rate catches oscillation; estimator catches confident-but-wrong.
+        self.error_estimator = nn.Linear(self.config.hidden_size, 1)
+
+        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),
+            # SHREK: zeros = no previous prediction — first step gives flip_rate ≈ 1.0
+            # device=H_init.device ensures prev_pred is on CUDA, matching z_H and logits
+            prev_pred=torch.zeros(batch_size, self.config.seq_len, dtype=torch.int32, device=self.H_init.device),
+            prev_q_halt=torch.full((batch_size,), -5.0, device=self.H_init.device),
+            prev_q_continue=torch.full((batch_size,), -5.0, device=self.H_init.device),
+        )
+        
+    def reset_carry(self, reset_flag: torch.Tensor, carry: HierarchicalReasoningModel_ACTV1InnerCarry):
+        # SHREK: zero out prev_pred for reset sequences so they start fresh.
+        # a reset sequence is one that just halted — it will solve a new puzzle next.
+        # zeroing prev_pred means first step gives flip_rate ≈ 1.0 (maximum uncertainty).
+        new_prev_pred = torch.where(
+            reset_flag.view(-1, 1),
+            torch.zeros_like(carry.prev_pred),
+            carry.prev_pred
+        )
+        new_prev_q_halt     = torch.where(reset_flag, torch.full_like(carry.prev_q_halt,     -5.0), carry.prev_q_halt)
+        new_prev_q_continue = torch.where(reset_flag, torch.full_like(carry.prev_q_continue, -5.0), carry.prev_q_continue)
+
+        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),
+            prev_pred=new_prev_pred,
+            prev_q_halt=new_prev_q_halt,
+            prev_q_continue=new_prev_q_continue,
+        )
+
+
+    # SHREK: removed task_type parameter — error signal is now universal (no task rules needed)
+    def forward(self, carry: HierarchicalReasoningModel_ACTV1InnerCarry, batch: Dict[str, torch.Tensor], require_trace=False):
+    #  -> Tuple[HierarchicalReasoningModel_ACTV1InnerCarry, torch.Tensor, Tuple[torch.Tensor, 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 — decode z_H into token predictions
+        output = self.lm_head(z_H)[:, self.puzzle_emb_len:]   # (B, seq_len, vocab_size)
+
+        # SHREK Component 1: Combined Error Signal
+        # Signal A — flip rate: what fraction of tokens changed from last step?
+        #   catches oscillation (model changing its mind between wrong options)
+        flip_err, current_pred = get_error_signal(output, carry.prev_pred)     # (B,), (B, seq_len)
+
+        # Signal B — learned estimator: reads z_H and predicts how wrong the model is.
+        #   catches stuck-but-wrong (model confidently on wrong answer without oscillating)
+        #   average over content positions (skip puzzle embedding prefix positions)
+        #   CRITICAL: detach z_H_mean so aux_loss only trains error_estimator weights,
+        #   NOT the main reasoning layers. Without detach, aux_loss sends parasitic
+        #   gradients through z_H → H_level/L_level that conflict with lm_loss.
+        z_H_mean    = z_H[:, self.puzzle_emb_len:].mean(dim=1).detach()        # (B, hidden_size)
+        learned_err = torch.sigmoid(self.error_estimator(z_H_mean.float()))    # (B, 1)
+        learned_err = learned_err.squeeze(-1)                                  # (B,)
+
+        # SHREK: combined error = 50/50 blend of both signals
+        # flip_err works immediately from step 1 (no learning needed)
+        # learned_err becomes accurate over training and takes over as the stronger signal
+        error = 0.5 * flip_err + 0.5 * learned_err                            # (B,)
+
+        # SHREK Component 2: Stagnation-Aware Q-head
+        # Ablation: can be disabled via config.enable_stagnation_delta=False
+        if self.config.enable_stagnation_delta:
+            # Compute how much z_H changed during this reasoning step.
+            # CRITICAL: detach so Q-head gradients don't flow through all positions of z_H.
+            # Without detach, Q-head loss would route gradients through all 82 positions,
+            # drowning out lm_loss. With detach, Q-head gradients only flow through position 0.
+            z_H_f = z_H.detach().float()
+            z_H_start = carry.z_H.detach().float()
+            delta = torch.norm(z_H_f - z_H_start, dim=(1, 2)) / (torch.norm(z_H_start, dim=(1, 2)) + 1e-6)  # (B,)
+            # Q-head reads CLS token + stagnation delta
+            q_input = torch.cat([z_H[:, 0].to(torch.float32), delta.unsqueeze(-1)], dim=-1)  # (B, hidden_size+1)
+        else:
+            # Fallback: same as original HRM — Q-head reads CLS token only
+            q_input = z_H[:, 0].to(torch.float32)  # (B, hidden_size)
+        q_logits = self.q_head(q_input).to(torch.float32)  # (B, 2)
+
+        # SHREK: inject combined error into z_H (AFTER Q-head, only affects carry for next step)
+        # Ablation: can be disabled via config.enable_error_injection=False
+        if self.config.enable_error_injection:
+            # error_encoder maps scalar -> hidden_size vector
+            # alpha follows linear warmup: 0 → alpha_max over warmup steps
+            # scaled by 1/sqrt(hidden_size) so injection is proportional to model size
+            error_emb = self.error_encoder(error.unsqueeze(-1))                    # (B, hidden_size)
+            # SHREK: compute alpha from warmup schedule (not learned)
+            # During warmup, alpha ramps linearly from 0 to alpha_max.
+            # After warmup, alpha stays at alpha_max.
+            # Uses torch.clamp instead of Python min() to stay compatible with torch.compile.
+            with torch.no_grad():
+                if self.training:
+                    self._alpha_step += 1
+                alpha = self.config.alpha_max * torch.clamp(self._alpha_step / self.config.alpha_warmup_steps, max=1.0)
+            scale = math.sqrt(self.config.hidden_size)
+            z_H = z_H + (alpha * error_emb.unsqueeze(1) / scale).to(z_H.dtype)   # (B, seq_len, hidden_size)
+
+        # New carry: store error-injected z_H so next ACT step starts from it
+        new_carry = HierarchicalReasoningModel_ACTV1InnerCarry(
+            z_H=z_H.detach(),
+            z_L=z_L.detach(),
+            # SHREK: store current predictions — next step compares against these for flip rate
+            prev_pred=current_pred.detach(),
+            # SHREK: Q-values written to carry for torch.compile compatibility (not read for Q-targets)
+            prev_q_halt=q_logits[..., 0].detach(),
+            prev_q_continue=q_logits[..., 1].detach(),
+        )
+
+        # SHREK: also return learned_err so pretrain.py can compute auxiliary loss
+        # auxiliary loss trains the estimator: predicted error should match real lm_loss
+        if require_trace:
+            return z_H_trace, new_carry, output, (q_logits[..., 0], q_logits[..., 1]), learned_err
+        else:
+            return new_carry, output, (q_logits[..., 0], q_logits[..., 1]), learned_err
+
+
+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()}
+        )
+        
+    # SHREK: removed task_type parameter — error signal is universal, no task rules needed
+    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()}
+
+        # SHREK: run inner forward — unpack learned_err for auxiliary loss in pretrain.py
+        if require_trace:
+            z_H_trace, new_inner_carry, logits, (q_halt_logits, q_continue_logits), learned_err = self.inner(new_inner_carry, new_current_data, require_trace=require_trace)
+        else:
+            new_inner_carry, logits, (q_halt_logits, q_continue_logits), learned_err = self.inner(new_inner_carry, new_current_data)
+
+        outputs = {
+            "logits": logits,
+            "q_halt_logits": q_halt_logits,
+            "q_continue_logits": q_continue_logits,
+            "learned_err": learned_err,   # SHREK: for auxiliary loss in pretrain.py
+        }
+        
+        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)
+
+                # SHREK: Q-target via double forward pass (same as original HRM).
+                # Run inner() a second time to get NEXT step's Q-values (step T+1).
+                # Save/restore _alpha_step so the second call doesn't double-count
+                # the warmup — without this, alpha reaches full strength at step ~2500
+                # instead of 5000, destabilizing early training.
+                # NOTE: prev_q fields in carry are not read here — they must stay
+                # in the dataclass for torch.compile compatibility.
+                saved_alpha_step = self.inner._alpha_step.clone()
+                next_q_halt_logits, next_q_continue_logits = self.inner(new_inner_carry, new_current_data)[-2]
+                self.inner._alpha_step.copy_(saved_alpha_step)
+
+                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/maze-hard/shrek-tiny/losses.py b/maze-hard/shrek-tiny/losses.py
new file mode 100644
index 0000000000000000000000000000000000000000..784535de2c64479d78806a5a3d48f24b96ce16e2
--- /dev/null
+++ b/maze-hard/shrek-tiny/losses.py
@@ -0,0 +1,130 @@
+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()
+
+        # SHREK: auxiliary loss — trains error_estimator to predict per-sample lm_loss.
+        # Computed here (not in pretrain.py) so learned_err still has gradients.
+        # per_sample_lm_loss is detached so only error_estimator weights get updated by aux_loss.
+        aux_loss = 0
+        if "learned_err" in outputs:
+            per_sample_lm_loss = (ds_mask * loss / loss_divisor).sum(dim=1).detach()  # (B,)
+            # SHREK: min-max normalize target to 0-1 so it matches sigmoid output range.
+            # When all losses are equal (early training), target = 0 (no distinguishable error).
+            # When losses vary, target spreads across 0-1 giving the estimator a useful signal.
+            lm_min = per_sample_lm_loss.min()
+            lm_max = per_sample_lm_loss.max()
+            target = (per_sample_lm_loss - lm_min) / (lm_max - lm_min + 1e-6)       # (B,) in [0, 1]
+            aux_loss = F.mse_loss(outputs["learned_err"], target, reduction="sum")
+            metrics["aux_loss"] = aux_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) + 0.1 * aux_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) + 0.1 * aux_loss, metrics, detached_outputs, new_carry.halted.all(), new_carry.halted & act_halt
diff --git a/maze-hard/shrek-tiny/step_101556 b/maze-hard/shrek-tiny/step_101556
new file mode 100644
index 0000000000000000000000000000000000000000..aa122bfbf195500277d6c7ce4b84fcd223298794
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_101556
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6712bc5973d62b335ed79937460ff83e358d6ddb6b822ddb7996603d98697d2e
+size 54578071
diff --git a/maze-hard/shrek-tiny/step_10416 b/maze-hard/shrek-tiny/step_10416
new file mode 100644
index 0000000000000000000000000000000000000000..dcfeeb94b3a0710acd9701fc0ca6a9989954ce17
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_10416
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:97847cfdfdd412d503ebd21d47ddcb7b4ab105df108bd0d06959f2f8c4511b23
+size 54578037
diff --git a/maze-hard/shrek-tiny/step_109368 b/maze-hard/shrek-tiny/step_109368
new file mode 100644
index 0000000000000000000000000000000000000000..09fc1f21b3ec6b47c74700c601ed9fc705f64f95
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_109368
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f1ad9e5b3acf54d63fb6235a233412e96c2bf85309598ecefae48f10b7e3c0c1
+size 54578071
diff --git a/maze-hard/shrek-tiny/step_11718 b/maze-hard/shrek-tiny/step_11718
new file mode 100644
index 0000000000000000000000000000000000000000..7837febb2476e99285df6c01762f3489b589f1ec
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_11718
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c151235d410419d56df7663518f692f4c6fbbb2a8228447c7f59299f9ce8a28b
+size 54578037
diff --git a/maze-hard/shrek-tiny/step_117180 b/maze-hard/shrek-tiny/step_117180
new file mode 100644
index 0000000000000000000000000000000000000000..04c672ceeba24ee2d0682b7ad6ebee7152e3bbe8
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_117180
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b0e5a544f8a20a9cce29844e6824cc1c85b0bbf255c601d81c2243037a176499
+size 54578071
diff --git a/maze-hard/shrek-tiny/step_124992 b/maze-hard/shrek-tiny/step_124992
new file mode 100644
index 0000000000000000000000000000000000000000..1cde54fab21052977c537659df271746cb0f96bb
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_124992
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:64af000661654565fd47466bec66d2af18df266128f8f89deba647b05a4eaf55
+size 54578071
diff --git a/maze-hard/shrek-tiny/step_1302 b/maze-hard/shrek-tiny/step_1302
new file mode 100644
index 0000000000000000000000000000000000000000..4b214ec1ae23e2291e157ddf22fd6ed5564e0475
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_1302
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:45afc1831f5e9f9aa00c72a3c5c9c1665c97fa8e6c13fde63dbf8713a24d82f6
+size 54578003
diff --git a/maze-hard/shrek-tiny/step_13020 b/maze-hard/shrek-tiny/step_13020
new file mode 100644
index 0000000000000000000000000000000000000000..d08042e171289f4c678ee9d0cf570aad0c839036
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_13020
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb3e76f765d5cee0e8085c5a33ecf69504082480238810cbdc021fc2ce2582a6
+size 54578037
diff --git a/maze-hard/shrek-tiny/step_132804 b/maze-hard/shrek-tiny/step_132804
new file mode 100644
index 0000000000000000000000000000000000000000..ec70f126156ac06b64b4c2690b004930e8f218f6
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_132804
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9b037d27b0ca9d7dd92afbc4d5bb5abab097b5161ec83a47f867dc805721ecba
+size 54578071
diff --git a/maze-hard/shrek-tiny/step_140616 b/maze-hard/shrek-tiny/step_140616
new file mode 100644
index 0000000000000000000000000000000000000000..dfcc2f4d3fec47242013b4b15de0ebc1927572aa
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_140616
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e28d1249563e68be3642ac067ac5e93855e3a6b95daa1425de562b4f0336b226
+size 54578071
diff --git a/maze-hard/shrek-tiny/step_14322 b/maze-hard/shrek-tiny/step_14322
new file mode 100644
index 0000000000000000000000000000000000000000..66a303a6d97c0188f9d3d8872a331509211c395a
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_14322
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f82db92c284aebee89bf5072b8d90ac9c296e7d89f7b049e221f0cd01dadcf8a
+size 54578037
diff --git a/maze-hard/shrek-tiny/step_148428 b/maze-hard/shrek-tiny/step_148428
new file mode 100644
index 0000000000000000000000000000000000000000..d0a97a40c74dc5e9e335a19c05d6c6fd444756f0
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_148428
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:151395c1e331bc1ca531f1fde3efcc1139a466e4168f62dda5bc0557281b2f72
+size 54578071
diff --git a/maze-hard/shrek-tiny/step_15624 b/maze-hard/shrek-tiny/step_15624
new file mode 100644
index 0000000000000000000000000000000000000000..a111de37c2b24ea51a4496846b6d3678cf41166c
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_15624
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:23da15f2b25e47733a1a854fd8d608292e543fe2cf2c91fc726cd7f2430bd451
+size 54578037
diff --git a/maze-hard/shrek-tiny/step_156240 b/maze-hard/shrek-tiny/step_156240
new file mode 100644
index 0000000000000000000000000000000000000000..82508dc19712156942f8a31cea088e10f3557a3e
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_156240
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:80ddf58803f8befa9c29a76285c79df11fb1574280997059f4b043beeb714669
+size 54578071
diff --git a/maze-hard/shrek-tiny/step_16926 b/maze-hard/shrek-tiny/step_16926
new file mode 100644
index 0000000000000000000000000000000000000000..6f5e9682444093926de991ea79bdd57e0b4305a4
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_16926
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5cf9eea19c38e39e6281b98101d8233ee96e2670081726dc8f291cc3de8c5232
+size 54578037
diff --git a/maze-hard/shrek-tiny/step_18228 b/maze-hard/shrek-tiny/step_18228
new file mode 100644
index 0000000000000000000000000000000000000000..ebac95d1da6d5ffad616c361f5683ee1323d1008
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_18228
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e006722481604267aac87d432e6636fea6ea0f45663aa39ed9b44debbcbb02ee
+size 54578037
diff --git a/maze-hard/shrek-tiny/step_19530 b/maze-hard/shrek-tiny/step_19530
new file mode 100644
index 0000000000000000000000000000000000000000..b6dc9d0bba1f08995f3445a9137816781f5f689c
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_19530
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d406fb4ebe3e30b7baf4f3aee1ffe3baab90429578205b5b2a70ce4d8efb527
+size 54578037
diff --git a/maze-hard/shrek-tiny/step_20832 b/maze-hard/shrek-tiny/step_20832
new file mode 100644
index 0000000000000000000000000000000000000000..c3eb8c5bb0e61c0368f1c75ad358e5134c3ad54f
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_20832
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2cdedd04177eed270120bfd1f9190ae5f8b7db8fa293c49bafc1193100cd0794
+size 54578037
diff --git a/maze-hard/shrek-tiny/step_22134 b/maze-hard/shrek-tiny/step_22134
new file mode 100644
index 0000000000000000000000000000000000000000..8b7bf37de92e682f0d4d3e93c4f91cd14161aeac
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_22134
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a278cfb1ef6c1d3884446f0e90ef9abfe303460647ce355ddbbd48aacad96ced
+size 54578037
diff --git a/maze-hard/shrek-tiny/step_23436 b/maze-hard/shrek-tiny/step_23436
new file mode 100644
index 0000000000000000000000000000000000000000..37db81d648463446815469ca09c87555c42bfe20
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_23436
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb5ed13c96bfe920f78608037d3964ed58bc914626cc84c9bbaaf88ab5ddd25c
+size 54578037
diff --git a/maze-hard/shrek-tiny/step_24738 b/maze-hard/shrek-tiny/step_24738
new file mode 100644
index 0000000000000000000000000000000000000000..2512bf5af895228eac17dcd2e31921a93940fcc1
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_24738
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b84e32d50884953cec18235dbc1d8124725f1eedbda9a395b02b1121291834b1
+size 54578037
diff --git a/maze-hard/shrek-tiny/step_2604 b/maze-hard/shrek-tiny/step_2604
new file mode 100644
index 0000000000000000000000000000000000000000..fe3a51eef11c248013b6f78f8f3c569c4ad183c5
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_2604
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6b88052cdb0c23704f5ee6bbd06fa1aefb703162c0342145cc28f8df7b11ddfc
+size 54578003
diff --git a/maze-hard/shrek-tiny/step_26040 b/maze-hard/shrek-tiny/step_26040
new file mode 100644
index 0000000000000000000000000000000000000000..e68644f5e887648a359da907f728dca87c6c44a2
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_26040
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d83be8741f65f508221bdc09b4e26c43c970991303a3938b7db39b85bd248922
+size 54578037
diff --git a/maze-hard/shrek-tiny/step_31248 b/maze-hard/shrek-tiny/step_31248
new file mode 100644
index 0000000000000000000000000000000000000000..e4e3e6ce61150115bfb3d166f8c69517992f5ae5
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_31248
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fd959b90d334cc643bed58c0dc6aaa2d75ec52c4416e99e601f0e650245f7d65
+size 54578037
diff --git a/maze-hard/shrek-tiny/step_3906 b/maze-hard/shrek-tiny/step_3906
new file mode 100644
index 0000000000000000000000000000000000000000..41548f10d5a483b344bb836517d68f7947b51cdc
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_3906
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7596997caa6023f13f3e038d079cd16dd2969e146593f80e68cfe9b12ceedea6
+size 54578003
diff --git a/maze-hard/shrek-tiny/step_39060 b/maze-hard/shrek-tiny/step_39060
new file mode 100644
index 0000000000000000000000000000000000000000..87a4f64de3974e62e4cc35f1db113b5aa8b0f6f1
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_39060
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2983268545303d6ceab2a35fd1bdb344db628926e742fe394d66dc1f8e5f7033
+size 54578037
diff --git a/maze-hard/shrek-tiny/step_46872 b/maze-hard/shrek-tiny/step_46872
new file mode 100644
index 0000000000000000000000000000000000000000..22b2da8a418da37fdd401968553259365cb4f957
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_46872
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f76e7172714bc200938b4fe048dc0d31e37c869ed83a26224b0e9f6a37453e8f
+size 54578037
diff --git a/maze-hard/shrek-tiny/step_5208 b/maze-hard/shrek-tiny/step_5208
new file mode 100644
index 0000000000000000000000000000000000000000..3c5166dda09eea2a95d9f91978bd6c54fa83af8e
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_5208
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bfac56b5475906c27d1f2fe42e2d044ccec45bf66e8c9fb437b163cfe5ed6f65
+size 54578003
diff --git a/maze-hard/shrek-tiny/step_54684 b/maze-hard/shrek-tiny/step_54684
new file mode 100644
index 0000000000000000000000000000000000000000..d7f22876a06b862e5978c698537c7e55de82ed80
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_54684
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a80c0cee3c091c6fe6bd3d0b9fa3c369116b37b0c853412f7339b0686bef88c5
+size 54578037
diff --git a/maze-hard/shrek-tiny/step_62496 b/maze-hard/shrek-tiny/step_62496
new file mode 100644
index 0000000000000000000000000000000000000000..92b30e58eb672eea74b34badc8e99423a394a20e
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_62496
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c3fbcb5ae62b11abd95160376337d159207a8b13986f20455b56423737f08b0
+size 54578037
diff --git a/maze-hard/shrek-tiny/step_6510 b/maze-hard/shrek-tiny/step_6510
new file mode 100644
index 0000000000000000000000000000000000000000..8f880fc8433a1da8bfaf2cd98dec60bb3aba0f49
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_6510
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:376601c1511c39e6d547ee62169a6d6f42da8c78ed96593859beeda7c7c92b2a
+size 54578003
diff --git a/maze-hard/shrek-tiny/step_70308 b/maze-hard/shrek-tiny/step_70308
new file mode 100644
index 0000000000000000000000000000000000000000..9454602743ffe7c36b3aab927c577b91bd416e48
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_70308
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:53410bc4b48775a34f78c8225a23002dde6a17aa238fc707a79f791f062e7def
+size 54578037
diff --git a/maze-hard/shrek-tiny/step_7812 b/maze-hard/shrek-tiny/step_7812
new file mode 100644
index 0000000000000000000000000000000000000000..c9f5a3d4bf911671a532da5d0aa1441d14b751c4
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_7812
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:54de7125620b0b16e4e5eea32e62d87a3a55e755d6b06f23a12505cfa3ced1b7
+size 54578003
diff --git a/maze-hard/shrek-tiny/step_78120 b/maze-hard/shrek-tiny/step_78120
new file mode 100644
index 0000000000000000000000000000000000000000..28ef577aed107bb26900300396f85d48b3efb229
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_78120
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bdf52d5f6208dd03e0a862dbbf081d3561ad80f56620c1c7d236aecd86f646f2
+size 54578037
diff --git a/maze-hard/shrek-tiny/step_85932 b/maze-hard/shrek-tiny/step_85932
new file mode 100644
index 0000000000000000000000000000000000000000..0ab537f49897d588b7de201f9571de332bcc7f73
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_85932
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:47347f1ec085968336af772e80f1870457c06cab294dea27037b25e6950eea32
+size 54578037
diff --git a/maze-hard/shrek-tiny/step_9114 b/maze-hard/shrek-tiny/step_9114
new file mode 100644
index 0000000000000000000000000000000000000000..b0921053c07843aeece619642bc0f1401893f477
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_9114
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c0d018e24f878a95112179eefa12b482ad8bc8a053ff203a059ba97bba0e5d0c
+size 54578003
diff --git a/maze-hard/shrek-tiny/step_93744 b/maze-hard/shrek-tiny/step_93744
new file mode 100644
index 0000000000000000000000000000000000000000..531d0b890331622201c8fbd78003af3facf5f228
--- /dev/null
+++ b/maze-hard/shrek-tiny/step_93744
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e6dbeea9752744c1d72c4ee3706ad29ea55c17781aa69beaf45a837f47e04752
+size 54578037