Layers of thought: smarter reasoning with the Hierarchical Reasoning Model

# Inside HRM.forward()
x_tilde = self.input_projector(x)  # (B, h_dim)

z_next = self.rnn_cell(input_combined, z_prev)
z_next = self.norm(z_next)  # RMSNorm keeps scale in check

l_input = torch.cat([x_tilde, zH], dim=-1)  # fuse current step with high-level intent
zL = self.l_module(zL, l_input)            # perform low-level update

h_input = torch.cat([zL, zH], dim=-1)  # what we just learned + prior plan
zH = self.h_module(zH, h_input)        # macro‑update

        # Project input into hidden reasoning space
        x_tilde = self.input_projector(x)  # (B, h_dim)

        # Initialize low-level and high-level memory states
        zL = self.l_module.init_state(batch_size, self.l_dim, self.device)
        zH = self.h_module.init_state(batch_size, self.h_dim, self.device)

        # N outer cycles (high-level reasoning updates)
        for n in range(self.n_cycles):
            # T low-level reasoning steps per cycle
            for t in range(self.t_steps):
                l_input = torch.cat([x_tilde, zH], dim=-1)  # (B, 2*h_dim)
                zL = self.l_module(zL, l_input)             # update zL

            # After T low-level steps, update high-level memory
            h_input = torch.cat([zL, zH], dim=-1)          # (B, l_dim + h_dim)
            zH = self.h_module(zH, h_input)                # update zH

        # Final prediction from abstract reasoning memory
        y_hat = self.output_projector(zH)                  # (B, output_dim)

class RMSNorm(nn.Module):
    """
    Root Mean Square Normalization.
    Normalizes across features while preserving scale via a learned weight.
    Used throughout HRM instead of LayerNorm.
    """
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def _norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        output = self._norm(x.float()).type_as(x)
        return output * self.weight

class RecurrentBlock(nn.Module):
    """
    A recurrent update block used by both L and H modules.
    Internally uses a GRUCell + RMSNorm for stable updates.
    """
    def __init__(self, input_dim, hidden_dim, name="RecurrentBlock"):
        super().__init__()
        self.name = name
        self.rnn_cell = nn.GRUCell(input_dim, hidden_dim)
        self.norm = RMSNorm(hidden_dim)

    def forward(self, z_prev, input_combined):
        """
        Forward step of the RNN.
        - z_prev: previous hidden state (B, hidden_dim)
        - input_combined: input at this step (B, input_dim)
        Returns: next hidden state (B, hidden_dim)
        """
        z_next = self.rnn_cell(input_combined, z_prev)
        z_next = self.norm(z_next)
        return z_next

    def init_state(self, batch_size, hidden_dim, device):
        """Returns a zero-initialized state."""
        return torch.zeros(batch_size, hidden_dim, device=device)

class InputProjector(nn.Module):
    """
    Projects the input embedding into the HRM hidden space.
    This is the 'x_tilde' used throughout reasoning.
    """
    def __init__(self, input_dim, hidden_dim):
        super().__init__()
        self.project = nn.Linear(input_dim, hidden_dim)
        self.norm = RMSNorm(hidden_dim)

    def forward(self, x):
        x_proj = self.project(x)
        x_tilde = self.norm(x_proj)
        return x_tilde

class OutputProjector(nn.Module):
    """
    Projects the final high-level hidden state (zH) to the output space.
    For HRM this is typically a scalar quality score.
    """
    def __init__(self, h_dim, output_dim):
        super().__init__()
        self.project = nn.Linear(h_dim, output_dim)

    def forward(self, zH_final):
        return self.project(zH_final)

class HRMModel(nn.Module):
    """
    Hierarchical Reasoning Model (HRM)

    Models layered reasoning using two coupled RNNs:
    - Low-level module (L): simulates fine-grained steps (e.g. CoT steps)
    - High-level module (H): aggregates abstract strategic updates

    The model processes reasoning traces through N nested cycles,
    each composed of T low-level updates and a single high-level update.
    """
    def __init__(self, cfg, logger=None):
        super().__init__()
        self.logger = logger

        # Model hyperparameters from config
        self.input_dim = cfg.get("hrm.input_dim", 2048)
        self.h_dim = cfg.get("hrm.h_dim", 256)
        self.l_dim = cfg.get("hrm.l_dim", 128)
        self.output_dim = cfg.get("hrm.output_dim", 1)
        self.n_cycles = cfg.get("hrm.n_cycles", 4)  # Outer loop depth
        self.t_steps = cfg.get("hrm.t_steps", 4)    # Inner loop steps

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # Input projection network
        self.input_projector = InputProjector(self.input_dim, self.h_dim)

        # Low-level module (L): operates on [x_tilde, zH] → updates zL
        self.l_module = RecurrentBlock(2 * self.h_dim, self.l_dim, name="LModule")

        # High-level module (H): operates on [zL, zH] → updates zH
        self.h_module = RecurrentBlock(self.l_dim + self.h_dim, self.h_dim, name="HModule")

        # Output layer from final zH
        self.output_projector = OutputProjector(self.h_dim, self.output_dim)

    def forward(self, x):
        """
        Executes the full HRM reasoning process.

        Args:
            x: Input tensor of shape (B, input_dim) typically a plan embedding
        Returns:
            y_hat: Final prediction (B, output_dim)
            intermediate_states: Final zL and zH for optional introspection
        """
        batch_size = x.size(0)

        # Project input into hidden reasoning space
        x_tilde = self.input_projector(x)  # (B, h_dim)

        # Initialize low-level and high-level memory states
        zL = self.l_module.init_state(batch_size, self.l_dim, self.device)
        zH = self.h_module.init_state(batch_size, self.h_dim, self.device)

        # N outer cycles (high-level reasoning updates)
        for n in range(self.n_cycles):
            # T low-level reasoning steps per cycle
            for t in range(self.t_steps):
                l_input = torch.cat([x_tilde, zH], dim=-1)  # (B, 2*h_dim)
                zL = self.l_module(zL, l_input)             # update zL

            # After T low-level steps, update high-level memory
            h_input = torch.cat([zL, zH], dim=-1)          # (B, l_dim + h_dim)
            zH = self.h_module(zH, h_input)                # update zH

        # Final prediction from abstract reasoning memory
        y_hat = self.output_projector(zH)                  # (B, output_dim)

        # Return prediction and final latent states (optional for training/debug)
        intermediate_states = {'zL_final': zL, 'zH_final': zH}
        return y_hat, intermediate_states

    def to(self, device):
        """
        Custom `.to()` to move internal state tracking.
        """
        super().to(device)
        self.device = device
        return self

for n in range(n_cycles):        # High-level reasoning cycles
    for t in range(t_steps):     # Low-level steps per cycle
        zL = LModule(zL, [x, zH])        # Update low-level thoughts
    zH = HModule(zH, [zL, zH])           # Adjust strategy

class HRMTrainer(BaseTrainer): 
    """
    Trainer Agent for the Hierarchical Reasoning Model (HRM).
    Integrates with Stephanie's training framework.
    """
    def __init__(self, cfg, memory=None, logger=None):
        super().__init__(cfg, memory, logger)
        
        # --- HRM Specific Config ---
        self.model_type = "hrm"
        self.embedding_type = memory.embedding.type
        embedding_dim = memory.embedding.dim
        self.input_dim = embedding_dim * 2
        self.h_dim = cfg.get("hrm.h_dim", 256)
        self.l_dim = cfg.get("hrm.l_dim", 128)
        self.output_dim = cfg.get("hrm.output_dim", 1) # 1 for score prediction
        self.n_cycles = cfg.get("hrm.n_cycles", 4)
        self.t_steps = cfg.get("hrm.t_steps", 4)
        self.lr = cfg.get("hrm.lr", 1e-4)
        self.epochs = cfg.get("hrm.epochs", 10)
        self.batch_size = cfg.get("hrm.batch_size", 32)

        # Device setup (inherited or set)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # Initialize the HRM model
        hrm_cfg = {
            "hrm.input_dim": self.input_dim,
            "hrm.h_dim": self.h_dim,
            "hrm.l_dim": self.l_dim,
            "hrm.output_dim": self.output_dim,
            "hrm.n_cycles": self.n_cycles,
            "hrm.t_steps": self.t_steps,
        }
        self.hrm_model = HRMModel(hrm_cfg, logger=self.logger).to(self.device)
        
        # Optimizer (AdamW as recommended)
        self.optimizer = AdamW(self.hrm_model.parameters(), lr=self.lr)
        
        # Loss function (MSE for regression, e.g., predicting a score)
        # Can be made configurable (e.g., CrossEntropy for classification)
        self.criterion = nn.MSELoss() 

        self.logger.log("HRMTrainerInitialized", {
            "model_type": self.model_type,
            "input_dim": self.input_dim,
            "h_dim": self.h_dim,
            "l_dim": self.l_dim,
            "output_dim": self.output_dim,
            "n_cycles": self.n_cycles,
            "t_steps": self.t_steps,
            "lr": self.lr,
            "device": str(self.device)
        })

    def train (self, samples, dimension) -> dict:
        self.logger.log("HRMTrainingStarted", {"epochs": self.epochs})

        dataloader = self._create_dataloader(samples)
        if dataloader is None:
             self.logger.log("HRMTrainingError", {"message": "Dataloader creation failed or insufficient samples."})
             return {"status": "failed", "message": "Dataloader creation failed."}

        # 2. Training Loop
        for epoch in range(self.epochs):
            epoch_loss = 0.0
            num_batches = 0
            for _, (x_batch, y_batch) in enumerate(dataloader):
                # Move data to device
                x_batch = x_batch.to(self.device)
                y_batch = y_batch.to(self.device)

                # Zero gradients
                self.optimizer.zero_grad()

                # Forward pass
                y_pred, intermediate_states = self.hrm_model(x_batch) # (B, output_dim)

                # Compute loss
                # Ensure y_batch has the correct shape for the loss function
                # e.g., if output_dim=1, y_batch should be (B, 1) or (B,)
                # MSELoss expects same shape for pred and target
                loss = self.criterion(y_pred, y_batch)

                # Backward pass (One-step gradient approximation)
                # PyTorch's autograd handles this naturally for the looped architecture
                # as long as we don't unroll the entire N*T steps explicitly in the graph
                # and use the final loss.
                loss.backward()

                # Update parameters
                self.optimizer.step()
                
                epoch_loss += loss.item()
                num_batches += 1
                
                # Optional: Log batch loss
                # self.logger.log("HRMTrainingBatch", {"epoch": epoch, "batch": batch_idx, "loss": loss.item()})

            # Log average epoch loss
            avg_epoch_loss = epoch_loss / num_batches if num_batches > 0 else 0.0
            self.logger.log("HRMTrainingEpoch", {"epoch": epoch, "avg_loss": avg_epoch_loss})

        # 3. Save Model
        self._save_model(dimension)
        
        self.logger.log("HRMTrainingCompleted", {"final_avg_loss": avg_epoch_loss})
        return {"status": "trained", "final_loss": avg_epoch_loss}

    def _create_dataloader(self, samples):
        """
        Creates a DataLoader for HRM training.
        Assumes samples contain context_text, document_text, and a target_score.
        This is a basic example. You might need more complex logic based on your
        specific task (e.g., predicting next step in a sequence).
        """
        valid_samples = []
        for s in samples:
            ctx_text = s.get("context_text", "") # Or goal_text
            doc_text = s.get("document_text", "") # Or scorable.text
            # Target for HRM training. This is crucial.
            # Example: Predicting a score (like SICQL Q-value) or a derived metric.
            target_value = s.get("target_score", s.get("score", None)) 
            
            # Example: Using SICQL score as target
            # target_value = s.get("sicql_q_value", None) 

            if not ctx_text or not doc_text or target_value is None:
                continue # Skip invalid samples

            try:
                ctx_emb = torch.tensor(self.memory.embedding.get_or_create(ctx_text), dtype=torch.float32)
                doc_emb = torch.tensor(self.memory.embedding.get_or_create(doc_text), dtype=torch.float32)
                target_tensor = torch.tensor([target_value], dtype=torch.float32) # Shape (1,) for MSE with output_dim=1
                
                # Input to HRM: Concatenated embeddings
                input_tensor = torch.cat([ctx_emb, doc_emb], dim=-1) # Shape (input_dim,)
                
                valid_samples.append((input_tensor, target_tensor))
            except Exception as e:
                self.logger.log("HRMDataError", {"error": str(e), "sample_id": s.get("id", "unknown")})
                continue

        if len(valid_samples) < self.min_samples: # Assuming min_samples is in cfg or BaseTrainer
            self.logger.log("HRMDataError", {"message": f"Insufficient valid samples: {len(valid_samples)} < {self.min_samples}"})
            return None

        # Create TensorDataset and DataLoader
        inputs, targets = zip(*valid_samples)
        dataset = TensorDataset(torch.stack(inputs), torch.stack(targets))
        dataloader = DataLoader(dataset, batch_size=self.batch_size, shuffle=True)
        self.logger.log("HRMDataLoaderCreated", {"num_samples": len(valid_samples), "num_batches": len(dataloader)})
        return dataloader

    def _save_model(self, dimension: str):
        """Saves the trained HRM model components using the Locator."""
        locator = self.get_locator(dimension) # Assuming BaseTrainer provides this
        
        # Save model state dict
        torch.save(self.hrm_model.state_dict(), locator.model_file(suffix="_hrm.pt"))
        
        # Save individual components if needed (optional, but matches SICQL pattern)
        # torch.save(self.hrm_model.input_projector.state_dict(), locator.model_file(suffix="_input.pt"))
        # torch.save(self.hrm_model.l_module.state_dict(), locator.model_file(suffix="_l.pt"))
        # torch.save(self.hrm_model.h_module.state_dict(), locator.model_file(suffix="_h.pt"))
        # torch.save(self.hrm_model.output_projector.state_dict(), locator.model_file(suffix="_output.pt"))
        
        # Save configuration
        meta = {
            "model_type": self.model_type,
            "input_dim": self.input_dim,
            "h_dim": self.h_dim,
            "l_dim": self.l_dim,
            "output_dim": self.output_dim,
            "n_cycles": self.n_cycles,
            "t_steps": self.t_steps,
            "lr": self.lr,
            "epochs": self.epochs,
        }
        self._save_meta_file(meta, dimension) # Assuming this method exists in BaseTrainer
        
        self.logger.log("HRMModelSaved", {"path": locator.base_path})

class HRMTrainerAgent(BaseAgent):
    """
    Agent to train the Hierarchical Reasoning Model (HRM) for multiple dimensions.
    Uses SICQL Q-values as training targets for each goal/document pair.
    """
    def __init__(self, cfg, memory=None, logger=None):
        super().__init__(cfg, memory, logger)
        self.dimensions = cfg.get("dimensions", [])  # e.g., ["alignment", "relevance"]

        self.trainer = HRMTrainer(cfg.get("hrm", {}), memory, logger)
        self.scorer = SICQLScorer(cfg.get("sicql", {}), memory, logger)


    async def run(self, context: dict) -> dict:
        goal = context.get("goal", {})
        goal_text = goal.get("goal_text", "")
        documents = context.get(self.input_key, [])

        if not documents:
            self.logger.log("HRMTrainingAgentError", {
                "message": "No documents provided for training.",
                "input_key": self.input_key
            })
            context[self.output_key] = {"status": "failed", "reason": "no documents"}
            return context

        dimensional_training_samples = {dim: [] for dim in self.dimensions}

        for doc in documents:
            try:
                scorable = ScorableFactory.from_dict(doc, TargetType.DOCUMENT)

                score_bundle = self.scorer.score(
                    goal=goal,
                    scorable=scorable,
                    dimensions=self.dimensions
                )

                for dimension in self.dimensions:
                    score_result = score_bundle.results.get(dimension)
                    if not score_result or score_result.q_value is None:
                        self.logger.log("HRMTrainingAgentWarning", {
                            "message": f"Missing q_value for dimension '{dimension}'",
                            "doc_id": scorable.id
                        })
                        continue

                    dimensional_training_samples[dimension].append({
                        "context_text": goal_text,
                        "document_text": scorable.text,
                        "target_score": score_result.q_value
                    })

            except Exception as e:
                self.logger.log("HRMTrainingAgentDataError", {
                    "message": "Error processing document.",
                    "doc_id": doc.get("id", "unknown"),
                    "error": str(e)
                })

        # Log how many samples were prepared
        for dim, samples in dimensional_training_samples.items():
            self.logger.log("HRMTrainingDataPrepared", {
                "dimension": dim,
                "num_samples": len(samples)
            })

        # Train the HRM per dimension
        training_results = {}
        try:
            for dimension, samples in dimensional_training_samples.items():
                if not samples:
                    training_results[dimension] = {"status": "skipped", "reason": "no samples"}
                    continue

                result = self.trainer.train(samples=samples, dimension=dimension)
                training_results[dimension] = result

                self.logger.log("HRMTrainingAgentCompleted", {
                    "dimension": dimension,
                    "result": result
                })

            # Update context with structured results
            context[self.output_key] = {
                "status": "completed",
                "dimensions": self.dimensions,
                "results": training_results,
            }

        except Exception as e:
            self.logger.log("HRMTrainingAgentError", {
                "message": "Error during HRM training execution.",
                "error": str(e)
            })
            context[self.output_key] = {
                "status": "failed",
                "message": str(e)
            }

        return context

class HRMScorer(BaseScorer):
    """
    Scorer that uses a trained Hierarchical Reasoning Model (HRM) to evaluate
    goal/document pairs. The HRM performs internal multi-step reasoning to
    produce a quality score.
    """
    def __init__(self, cfg, memory, logger):
        super().__init__(cfg, memory, logger)
        self.model_type = "hrm" # This identifies the scorer type
        
        # Use the embedding details from memory
        self.embedding_type = self.memory.embedding.type
        self.dim = self.memory.embedding.dim
        # HRM might use a different internal dimension (h_dim), but input is based on self.dim
        # h_dim, l_dim, etc. are loaded from the model's meta file or config
        
        # Get target type and version from config, with defaults
        self.target_type = cfg.get("target_type", "document")
        self.model_path = cfg.get("model_path", "models")
        self.version = cfg.get("model_version", "v1")

        # The specific HRM task/dimension this scorer represents
        # This should match the `hrm_dimension` used during training
        self.hrm_dimension = cfg.get("hrm_dimension", "sicql_alignment") 
        
        # Dictionary to hold the loaded HRM model instance
        self.model = None
        # Dictionary to hold model metadata (e.g., hyperparameters)
        self.model_meta = None

        # Attempt to load the model during initialization
        self._load_model()

    def _load_model(self):
        """
        Loads the trained HRM model components and metadata using ModelLocator.
        """
        try:
            # Use the inherited get_locator method (from ModelLocatorMixin via BaseScorer)
            # This will create the path based on embedding_type, model_type (hrm), 
            # target_type, dimension (hrm_dimension), and version.
            locator = self.get_locator(self.hrm_dimension) 

            # Check if the model files exist Is right that is wrong
            model_file_path = locator.model_file(suffix="_hrm.pt") # Match the suffix used in saving
            meta_file_path = locator.meta_file()

            if not os.path.exists(model_file_path):
                self.logger.log("HRMScorerModelError", {
                    "message": "HRM model file not found.",
                    "path": model_file_path,
                    "dimension": self.hrm_dimension
                })
                return # Cannot load if file is missing

            # Load model metadata
            if os.path.exists(meta_file_path):
                self.model_meta = load_json(meta_file_path)
                self.logger.log("HRMScorerMetaLoaded", {
                    "dimension": self.hrm_dimension,
                    "meta": self.model_meta # Log key meta info if needed
                })
            else:
                self.logger.log("HRMScorerWarning", {
                    "message": "HRM meta file not found. Using defaults.",
                    "path": meta_file_path
                })
                self.model_meta = {} # Use empty dict if meta is missing

            # --- Reconstruct HRM Model Configuration ---
            # Get HRM hyperparameters from meta or use defaults consistent with training
            hrm_cfg_from_meta = {
                "hrm.input_dim": self.model_meta.get("input_dim", self.dim * 2), # Default concat
                "hrm.h_dim": self.model_meta.get("h_dim", 256),
                "hrm.l_dim": self.model_meta.get("l_dim", 128),
                "hrm.output_dim": self.model_meta.get("output_dim", 1),
                "hrm.n_cycles": self.model_meta.get("n_cycles", 4),
                "hrm.t_steps": self.model_meta.get("t_steps", 4),
                # lr, epochs are not needed for inference
            }
            
            # --- Instantiate HRM Model ---
            # Create an instance of the HRMModel with the loaded config
            self.model = HRMModel(hrm_cfg_from_meta, logger=self.logger)
            
            # --- Load Model Weights ---
            # Load the saved state dictionary into the model instance
            # Make sure the device is consistent
            self.model.to(self.device)
            self.model.load_state_dict(torch.load(model_file_path, map_location=self.device))
            self.model.eval() # Set to evaluation mode
            
            self.logger.log("HRMScorerModelLoaded", {
                "dimension": self.hrm_dimension,
                "model_path": model_file_path,
                "device": str(self.device)
            })

        except Exception as e:
            self.logger.log("HRMScorerInitError", {
                "message": "Failed to load HRM model.",
                "dimension": self.hrm_dimension,
                "error": str(e)
            })
            self.model = None # Ensure model is None on failure

    def score(self, goal: dict, scorable: Scorable, dimensions: list[str]) -> ScoreBundle:
        """
        Scores a single scorable item against a goal using the trained HRM model.
        
        Args:
            goal: A dictionary containing goal information (e.g., {"goal_text": "..."})
            scorable: A Scorable object representing the item to be scored.
            dimensions: A list of dimension names. The HRM scorer typically
                        produces one primary score, but this list allows integration
                        into the standard scoring framework. It will score if 
                        self.hrm_dimension is in this list.
        
        Returns:
            ScoreBundle: Contains the HRM score result if applicable.
        """
        results = {}

        if not self.model:
             self.logger.log("HRMScorerError", {
                 "message": "HRM model not loaded. Cannot score.",
                 "dimension": self.hrm_dimension
             })
             return ScoreBundle(results={})

        try:
            goal_text = goal.get("goal_text", "")
            doc_text = scorable.text

            if not goal_text or not doc_text:
                self.logger.log("HRMScorerWarning", {
                    "message": "Missing goal_text or scorable text.",
                    "dimension": self.hrm_dimension
                })
                return ScoreBundle(results={})

            # 1. Get embeddings
            ctx_emb_np = self.memory.embedding.get_or_create(goal_text)
            doc_emb_np = self.memory.embedding.get_or_create(doc_text)

            # 2. Convert to PyTorch tensors and move to device
            ctx_emb = torch.tensor(ctx_emb_np, dtype=torch.float32).to(self.device).unsqueeze(0)
            doc_emb = torch.tensor(doc_emb_np, dtype=torch.float32).to(self.device).unsqueeze(0)

            # 3. Prepare input for HRM Model (concatenate)
            x_input = torch.cat([ctx_emb, doc_emb], dim=-1) # Shape: (1, input_dim)

            # 4. Run the HRM Model (in evaluation mode) - Capture intermediate states
            with torch.no_grad():
                # UNPACK the tuple returned by HRMModel.forward
                # y_pred is the output tensor, intermediate_states is the dict
                y_pred, intermediate_states = self.model(x_input) # Shapes: (1, 1), dict
            
            # 5. Extract the scalar score value
            raw_hrm_score = y_pred.squeeze().item()

            # 6. Process intermediate states for logging/rationale
            # Extract final states (they are tensors)
            zL_final_tensor = intermediate_states.get('zL_final')
            zH_final_tensor = intermediate_states.get('zH_final')

            # Example: Calculate magnitude (L2 norm) of final states as a simple metric
            zL_magnitude = None
            zH_magnitude = None
            if zL_final_tensor is not None:
                # .item() to get scalar value from single-element tensor
                zL_magnitude = torch.norm(zL_final_tensor, p=2).item() 
            if zH_final_tensor is not None:
                zH_magnitude = torch.norm(zH_final_tensor, p=2).item()

            # Example: Get the actual final hidden state values (useful for debugging small models)
            # Convert to list for JSON serialization if needed
            # zL_final_values = zL_final_tensor.flatten().tolist() if zL_final_tensor is not None else None
            # zH_final_values = zH_final_tensor.flatten().tolist() if zH_final_tensor is not None else None

            # 7. (Optional) Apply post-processing/clipping/normalization
            final_score = raw_hrm_score # Or apply clipping/transform

            # 8. Create ScoreResult with enhanced rationale and metadata
            prompt_hash = ScoreORM.compute_prompt_hash(goal_text, scorable)

            # Build a more detailed rationale using intermediate state info
            rationale_parts = [f"HRM prediction (raw={round(raw_hrm_score, 4)})"]
            if zL_magnitude is not None:
                rationale_parts.append(f"zL_mag={round(zL_magnitude, 4)}")
            if zH_magnitude is not None:
                rationale_parts.append(f"zH_mag={round(zH_magnitude, 4)}")
            rationale = f" after {self.model_meta.get('n_cycles', 'N')}/{self.model_meta.get('t_steps', 'T')} cycles/steps. " + ", ".join(rationale_parts)

            # Prepare extra metadata to store in ScoreResult (optional)
            # This could include the magnitudes or even the full state lists (if small/serializable)
            extra_metadata = {
                "hrm_zL_final_magnitude": zL_magnitude,
                "hrm_zH_final_magnitude": zH_magnitude,
                # "hrm_zL_final_values": zL_final_values, # Uncomment if storing full states
                # "hrm_zH_final_values": zH_final_values, # Uncomment if storing full states
                "hrm_cycles": self.model_meta.get('n_cycles'),
                "hrm_t_steps": self.model_meta.get('t_steps'),
            }

            score_result = ScoreResult(
                dimension=self.hrm_dimension,
                score=final_score,
                rationale=rationale, # Enhanced rationale
                weight=1.0,
                q_value=raw_hrm_score,
                energy=raw_hrm_score, # You might adjust this based on intermediate states if desired
                source=self.model_type,
                target_type=scorable.target_type,
                prompt_hash=prompt_hash,
            )

            # 8a. (Alternative) If ScoreResult can't hold extra metadata easily,
            # log the intermediate state info separately
            self.logger.log("HRMScorerIntermediateStates", {
                "dimension": self.hrm_dimension,
                "goal_id": goal.get("id", "unknown"),
                "scorable_id": scorable.id,
                "zL_final_magnitude": zL_magnitude,
                "zH_final_magnitude": zH_magnitude,
                # "zL_final_values": zL_final_values, # Log full values if needed/debugging
                # "zH_final_values": zH_final_values,
            })

            # 9. Add to results dictionary
            results[self.hrm_dimension] = score_result

            # 10. Log the scoring event
            self.logger.log("HRMScorerEvaluated", {
                "dimension": self.hrm_dimension,
                "goal_id": goal.get("id", "unknown"),
                "scorable_id": scorable.id,
                "raw_score": raw_hrm_score,
                "final_score": final_score,
                "zL_final_magnitude": zL_magnitude, # Log key metrics here too
                "zH_final_magnitude": zH_magnitude,
            })

        except Exception as e:
            self.logger.log("HRMScorerError", {
                "message": "Error during HRM scoring.",
                "dimension": self.hrm_dimension,
                "goal_id": goal.get("idHi Sime", "unknown"),
                "scorable_id": scorable.id,
                "error": str(e)
            })
            return ScoreBundle(results={})

        return ScoreBundle(results=results)


    def __repr__(self):
        return f"<HRMScorer(model_type={self.model_type}, dimension={self.hrm_dimension}, loaded={self.model is not None})>"

target_id,target_type,dimension,source,score,llm_score,delta
1,document,alignment,ebt,75.9662,70.0,5.966200000000001
1,document,alignment,hrm,9.906468391418457,70.0,-60.09353160858154
1,document,alignment,mrq,76.4469,70.0,6.446899999999999
1,document,alignment,sicql,100.0,70.0,30.0
1,document,alignment,svm,76.83712967511364,70.0,6.837129675113644
2,document,alignment,ebt,75.8436,65.0,10.843599999999995
2,document,alignment,hrm,9.910690307617188,65.0,-55.08930969238281
2,document,alignment,mrq,76.4469,65.0,11.4469
2,document,alignment,sicql,100.0,65.0,35.0

### Model Vs Llm Score Correlation
- **Source:** `hrm`
- **Dimension:** `alignment`
- **Description:** Correlation between model's raw score (from attributes) and LLM score.
- **Metric:** `Pearson Correlation Coefficient`
- **Value:** `0.2994345745505696`
- **P-Value:** `0.002474340154950069`
- **Sample Size:** `100`

@dataclass
class ExecutionStep:
    """
    Represents a single step in the execution of a reasoning plan.
    This can be generated by an executor like EpistemicPlanExecutorAgent.
    """
    step_id: Union[str, int]  # Unique identifier for the step (e.g., index, name)
    description: str  # A textual description of what this step does
    output_text: str  # The textual output or result of this step

    # The scores assigned to this step's output by various scorers (SICQL, EBT, etc.)
    # against the original goal. 
    scores: Optional[ScoreBundle] 

    plan_trace_id: Optional[int] = None  # Foreign key to the PlanTrace this step belongs to
    step_order: Optional[int] = None  # Foreign key to the PlanTrace this step belongs to
    # Optional: Embedding of the output_text. Can be computed on demand if not stored.
    
    # Optional: Any other metadata specific to this step
    extra_data: Optional[Dict[str, Any]] = field(default_factory=dict) 

    def to_dict(self) -> Dict[str, Any]:
        return {
            "step_id": self.step_id,
            "description": self.description,
            "output_text": self.output_text,
            "scores": self.scores.to_dict(),
            "plan_trace_id": self.plan_trace_id,
            "step_order": self.step_order,
            "extra_data": self.extra_data,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ExecutionStep":
        from stephanie.scoring.score_bundle import \
            ScoreBundle  # Local import to avoid circular dependencies

        return cls(
            step_id=data.get("step_id"),
            description=data.get("description", ""),
            output_text=data.get("output_text", ""),
            scores=ScoreBundle.from_dict(data.get("scores", {})),
            plan_trace_id=data.get("plan_trace_id"),
            step_order=data.get("step_order"),
            extra_data=data.get("extra_data", {}),
        )

@dataclass
class PlanTrace:
    """
    Represents the complete execution trace of a reasoning plan.
    This is the primary input for the EpistemicTraceEncoder and subsequently 
    the Epistemic Plan HRM model.
    """
    # --- Core Identifiers ---
    trace_id: str # Unique identifier for this specific trace/execution
    
    # --- Initial Context ---
    goal_text: str # The original goal or query
    goal_id: int
    input_data: Dict[str, Any] # Any initial data or variables provided to the plan
    
    # --- Plan Definition (Optional but useful for context) ---
    # This could be a representation of the DSPy program or pipeline used.
    # A simple string signature or a more structured representation.
    plan_signature: str 

    # --- Execution Details ---
    execution_steps: List[ExecutionStep] # The sequence of steps executed
    
    # --- Final Outcome ---
    final_output_text: str # The final output produced by the plan
    # The scores assigned to the final output by various scorers.
    final_scores: Optional[ScoreBundle] = None

    # --- Target for Epistemic Plan HRM Training ---
    # This is the label the HRM model will try to predict.
    # It represents the "epistemic quality" of this reasoning process.
    target_epistemic_quality: Optional[float] = None 
    # Source of the target quality score (e.g., "llm_judgment", "proxy_metric_avg_sicql_q")
    target_epistemic_quality_source: Optional[str] = None 

    # --- Metadata ---
    created_at: str = "" # ISO format timestamp
    # Any other execution metadata (e.g., time taken, DSPy optimizer version)
    extra_data: Optional[Dict[str, Any]] = field(default_factory=dict) 

    def to_dict(self) -> dict:
        return {
            "trace_id": self.trace_id,
            "goal_text": self.goal_text,
            "goal_id": self.goal_id,
            "input_data": self.input_data,
            "plan_signature": self.plan_signature,
            "execution_steps": [step.to_dict() for step in self.execution_steps],
            "final_output_text": self.final_output_text,
            "final_scores": self.final_scores.to_dict(),
            "target_epistemic_quality": self.target_epistemic_quality,
            "target_epistemic_quality_source": self.target_epistemic_quality_source,
            "created_at": self.created_at,
            "extra_data": self.extra_data,
        }

    def get_target_quality(self) -> float:
        if self.has_target_quality():
            return float(self.target_epistemic_quality)
        raise ValueError(f"Trace {self.trace_id} is missing 'target_epistemic_quality'")

    def has_target_quality(self) -> float:
        return self.target_epistemic_quality is not None

    # --- Utility Methods ---
    def get_all_text_outputs(self) -> List[str]:
        """Get a list of all text outputs, including intermediate steps and final output."""
        texts = [step.output_text for step in self.execution_steps]
        texts.append(self.final_output_text)
        return texts

    def get_all_score_bundles(self) -> List[ScoreBundle]:
        """Get a list of all ScoreBundles, including intermediate steps and final output."""
        bundles = [step.scores for step in self.execution_steps]
        bundles.append(self.final_scores)
        return bundles

    def to_markdown(self) -> str:
        lines = [f"## Plan Trace: {self.trace_id}", f"**Goal:** {self.goal_text}\n"]
        for step in self.execution_steps:
            step_id_str = str(step.step_id) if step.step_id is not None else "N/A"
            lines.append(f"### Step {step_id_str}: {step.description}")
            lines.append(f"Output: `{step.output_text}`")
            lines.append(step.scores.to_report(f"Step {step_id_str}: Scores"))
        lines.append(f"\n**Final Output:** `{self.final_output_text}`")
        lines.append("Final Scores:")
        lines.append(self.final_scores.to_report("Trace Final Scores") if self.final_scores else "No final scores available.")
        return "\n".join(lines)

    def save_as_markdown(self, reports_dir: str = "reports") -> str:
        os.makedirs(reports_dir, exist_ok=True)
        markdown_text = self.to_markdown()
        safe_trace_id = "".join(c for c in self.trace_id if c.isalnum() or c in (' ', '-', '_')).rstrip()
        filename = f"{safe_trace_id}.md"
        filepath = os.path.join(reports_dir, filename)
        with open(filepath, "w", encoding="utf-8") as f:
            f.write(markdown_text)
        return filepath

    def save_as_json(self, dir_path: str = "reports/json") -> str:
        os.makedirs(dir_path, exist_ok=True)
        filename = f"{self.trace_id}.json"
        path = os.path.join(dir_path, filename)
        with open(path, "w", encoding="utf-8") as f:
            json.dump(self.to_dict(), f, indent=2)

        print(f"PlanTraceSavedAsJSON path: {path}")

        return path

    @classmethod
    def from_dict(cls, data: dict) -> "PlanTrace":
        from stephanie.scoring.score_bundle import ScoreBundle

        execution_steps = [
            ExecutionStep(
                step_id=step["step_id"],
                description=step["description"],
                output_text=step["output_text"],
                scores=ScoreBundle.from_dict(step["scores"]),
                plan_trace_id=step.get("plan_trace_id"),
                step_order=step.get("step_order"),
                extra_data=step.get("extra_data", {}),
            )
            for step in data["execution_steps"]
        ]

        return cls(
            trace_id=data["trace_id"],
            goal_text=data["goal_text"],
            goal_id=data["goal_id"],
            input_data=data["input_data"],
            plan_signature=data["plan_signature"],
            execution_steps=execution_steps,
            final_output_text=data["final_output_text"],
            final_scores=ScoreBundle.from_dict(data["final_scores"]),
            target_epistemic_quality=data.get("target_epistemic_quality"),
            target_epistemic_quality_source=data.get("target_epistemic_quality_source"),
            created_at=data.get("created_at", ""),
            extra_data=data.get("extra_data", {}),
        )

# Define a Signature for a single LATS-style reasoning step
class ReasoningStepSignature(dspy.Signature):
    """Generate the next logical reasoning step towards solving a goal."""
    
    goal = dspy.InputField(desc="The main goal to solve.")
    previous_steps_summary = dspy.InputField(desc="A concise summary of the previous reasoning steps taken so far.")
    input_data = dspy.InputField(desc="Any initial data or context provided for the task.", format=lambda x: json.dumps(x, indent=2))
    step_number = dspy.InputField(desc="The current step number in the sequence.")
    
    # The output field instructs the model on the expected format
    next_step = dspy.OutputField(desc="The next reasoning step. Be specific and build on prior steps. "
                                      "If you have logically concluded the task, start your response EXACTLY with 'Final Answer: ' followed by your conclusion.")

FINAL_ANSWER_PATTERN = re.compile(r"(?:^|\n)\s*final\s*answer\s*[:：]\s*", re.IGNORECASE)


class EpistemicPlanExecutorAgent(BaseAgent):
    """
    Agent to execute a reasoning plan using a simplified, internal LATS-like process
    and generate a detailed PlanTrace for subsequent analysis by the Epistemic Plan HRM.
    This avoids direct dependency on the external LATSDSPyAgent.
    """

    def __init__(
        self, cfg: Dict[str, Any], memory: Any = None, logger: Any = None
    ):
        super().__init__(cfg, memory, logger)
        self.dimensions = cfg.get("dimensions", [])
        self.plan_timeout_seconds = cfg.get("plan_timeout_seconds", 300)
        self.max_reasoning_steps = cfg.get("max_reasoning_steps", 5) # Configurable steps
        self.use_hrm_in_trace = cfg.get("use_hrm_in_trace", True) # Config flag

        self.sicql_scorer = SICQLScorer(cfg=self.cfg.get("sicql", {}), memory=memory, logger=logger)
        if self.use_hrm_in_trace:
            self.hrm_scorer = HRMScorer(cfg=self.cfg.get("hrm", {}), memory=memory, logger=logger)
        else:
            self.hrm_scorer = None
        # Get the configured LM
        self.lm = dspy.LM(
            "ollama_chat/qwen3",
            api_base="http://localhost:11434",
            api_key="",
        )
        dspy.configure(lm=self.lm)
        self.step_predictor = dspy.ChainOfThought(
            signature=ReasoningStepSignature
        )
        self.logger.log("EpistemicPlanExecutorAgentInitialized", {
            "max_reasoning_steps": self.max_reasoning_steps,
            "use_hrm_in_trace": self.use_hrm_in_trace,
        })

    async def _run_simplified_lats(self, goal_text: str, input_data: Dict[str, Any]) -> List[str]:
        """
        Simplified internal logic to generate a sequence of reasoning steps,
        using dspy.Predict/ChainOfThought for structured prompting.

        Args:
            goal_text (str): The main goal to reason about.
            input_data (dict): Initial data provided to the reasoning process.

        Returns:
            List[str]: A list of strings, each representing an intermediate reasoning step/output.
        """
        trace_outputs = []
        # Start with an empty summary; the predictor can handle this.
        previous_steps_summary = "" 

        for step_num in range(1, self.max_reasoning_steps + 1):
            # self.logger.log("LATS_StepStarted", {"step": step_num, "summary": previous_steps_summary[-100:]})
            self.logger.log("LATS_StepStarted", {"step": step_num})

            try:
                # --- Use dspy.Predict/ChainOfThought to generate the next step ---
                # Prepare the prediction inputs based on the Signature
                prediction_kwargs = {
                    "goal": goal_text,
                    "previous_steps_summary": previous_steps_summary,
                    "input_data": input_data,
                    "step_number": step_num
                }
                
                prediction = self.step_predictor(**prediction_kwargs)
                
                # --- Extract the Output ---
                # The output is accessed via the attribute name defined in the Signature ('next_step')
                step_output_text = prediction.next_step.strip()

                # --- Check for Final Answer ---
                is_final_answer = bool(FINAL_ANSWER_PATTERN.search(step_output_text))

                is_final_answer = step_output_text.lower().startswith("final answer: ")
                if is_final_answer:
                    # Extract the part after "Final Answer: "
                    # final_part = step_output_text[len("final answer: "):].strip()
                    # trace_outputs.append(f"Final Answer: {final_part}")
                    # Let's keep the full text including the prefix for clarity in the trace
                    trace_outputs.append(step_output_text) 
                    self.logger.log("EpistemicPlanExecutorLATS", {
                        "message": f"Early stopping at step {step_num} due to 'Final Answer' signal.",
                        "final_answer_snippet": step_output_text[:100]
                    })
                    break # Stop the loop
                else:
                    trace_outputs.append(step_output_text)
                    # Update the summary for the next step
                    # A more robust summary could be built, but for now, append the last step
                    # Truncate previous summary and current step to keep it manageable
                    if len(previous_steps_summary) > 300:
                        previous_steps_summary = previous_steps_summary[-200:]
                    previous_steps_summary += f"\nStep {step_num}: {step_output_text[:100]}..."
                    # Ensure it doesn't grow too large
                    if len(previous_steps_summary) > 500:
                        previous_steps_summary = previous_steps_summary[-400:]

                self.logger.log("LATS_StepCompleted", {"step": step_num, "output_snippet": step_output_text[:100]})

            except Exception as e:
                self.logger.log("EpistemicPlanExecutorLATSStepError", {
                    "message": f"Error generating LATS-like step {step_num}.",
                    "error": str(e),
                    "traceback": traceback.format_exc(),
                })
                # Decide whether to break or continue with a placeholder/error step
                trace_outputs.append(f"[ERROR: Failed to generate step {step_num}]")
                # Continue to next step

        return trace_outputs

    async def run(self, context: Dict[str, Any]) -> Dict[str, Any]:

        existing_goal_ids = {
            pt.goal_id for pt in self.memory.plan_traces.all()
            if pt.goal_id is not None
        }
        goals = self.memory.goals.get_all_goals()

        for goal in goals:
            goal_id = goal.id
            if goal.id in existing_goal_ids:
                self.logger.log("EpistemicPlanExecutorSkipped", {
                    "goal_id": goal.id,
                    "message": "Goal already has a PlanTrace, skipping."
                })
                continue

            goal_dict = goal.to_dict()
            goal_text = goal.goal_text
            if not goal_text or len(goal_text) < 10:
                self.logger.log("EpistemicPlanExecutorWarning", {
                    "message": f"Goal text is too short or missing: {goal_text}",
                    "goal_id": goal.id
                })
                continue
            
            input_data = context.get("input_data", {})
            self.logger.log("EpistemicPlanExecutorStarted", {
                "goal_id": goal_id,
                "goal_text": goal_text,
                "input_data": input_data
            })

            if not goal_text:
                error_msg = "Missing 'goal_text' in context['goal']. Cannot execute plan."
                self.logger.log("EpistemicPlanExecutorError", {"message": error_msg})
                context[self.output_key] = {
                    "goal_id": goal_id,
                    "executor_agent": self.__class__.__name__,
                    "source": "simplified_lats_execution",
                    "status": "failed",
                    "error": error_msg
                }
                return context

            trace_id = f"trace_{uuid.uuid4().hex}"
            plan_signature = f"SimplifiedLATS_{self.max_reasoning_steps}_steps"

            execution_steps: List[ExecutionStep] = []
            final_output_text: str = ""
            final_scores: Optional[ScoreBundle] = None

            try:
                # --- Execute the Simplified LATS-like Reasoning ---
                trace_outputs = await self._run_simplified_lats(goal_text, input_data)

                # --- Process Generated Trace into ExecutionSteps ---
                step_id_counter = int(time.time() * 1000)
                processed_trace_info = []

                for i, step_output_text in enumerate(trace_outputs):
                    step_id_counter += 1
                    step_description = f"Simplified LATS Step {i + 1}"
                    processed_trace_info.append({
                        "step_id": step_id_counter,
                        "description": step_description,
                        "output_text": step_output_text.strip() # Clean up whitespace
                    })

                # --- Score Each Processed Step Using Stephanie Scorers ---
                for step_info in processed_trace_info:
                    step_id = step_info["step_id"]
                    step_description = step_info["description"]
                    step_output_text = step_info["output_text"]

                    if not step_output_text:
                        self.logger.log("EpistemicPlanExecutorWarning", {
                            "message": f"Generated step {step_id} has empty output. Skipping scoring."
                        })
                        continue

                    try:
                        scorable_dict = {"text": step_output_text, "id": str(step_id)} # Ensure ID is string
                        scorable = ScorableFactory.from_dict(scorable_dict, TargetType.DOCUMENT)

                        # --- Score the Step Output ---
                        sicql_scores: ScoreBundle = self.sicql_scorer.score(
                            goal=goal_dict, scorable=scorable, dimensions=self.dimensions
                        )
                        hrm_scores: Optional[ScoreBundle] = None
                        if self.hrm_scorer:
                            hrm_scores = self.hrm_scorer.score(
                                goal=goal_dict, scorable=scorable, dimensions=self.dimensions
                            )
                            if hrm_scores:
                                sicql_scores = sicql_scores.merge(hrm_scores)


                        # --- Create ExecutionStep Object ---
                        step_meta = {
                            "sicql_scores": sicql_scores.to_dict(),
                            "source": "simplified_lats_step"
                        }
                        if hrm_scores:
                            step_meta["hrm_scores"] = hrm_scores.to_dict()

                        exec_step = ExecutionStep(
                            step_id=str(step_id), # Ensure ID is string
                            description=step_description,
                            output_text=step_output_text,
                            scores=sicql_scores, # Primary scores for the trace
                            extra_data=step_meta,
                        )
                        execution_steps.append(exec_step)

                    except Exception as e:
                        self.logger.log("EpistemicPlanExecutorStepError", {
                            "message": f"Error scoring generated step {step_id}.",
                            "step_output_snippet": step_output_text[:50],
                            "error": str(e),
                            "traceback": traceback.format_exc(),
                        })
                        continue # Continue with other steps

                # --- Determine Final Output ---
                # The final output is typically the last step's text
                # Or, if the last step started with "Final Answer:", extract that part
                if execution_steps:
                    last_step_text = execution_steps[-1].output_text
                    if last_step_text.lower().startswith("final answer:"):
                        # Extract the part after "Final Answer:"
                        final_output_text = last_step_text[len("final answer:"):].strip()
                    else:
                        final_output_text = last_step_text
                else:
                    final_output_text = "No reasoning steps were generated."

                # --- Score the Final Output ---
                try:
                    final_scorable_dict = {"text": final_output_text, "id": f"{trace_id}_final"}
                    final_scorable = ScorableFactory.from_dict(final_scorable_dict, TargetType.DOCUMENT)
                    final_scores: ScoreBundle = self.sicql_scorer.score(
                        goal=goal_dict, scorable=final_scorable, dimensions=self.dimensions
                    )

                except Exception as e:
                    self.logger.log("EpistemicPlanExecutorFinalScoringError", {
                        "message": "Error scoring final output.",
                        "final_output_snippet": final_output_text[:50],
                        "error": str(e),
                        "traceback": traceback.format_exc(),
                    })

            except Exception as e:
                self.logger.log("EpistemicPlanExecutorExecutionError", {
                    "message": "Error during simplified LATS execution or trace processing.",
                    "error": str(e),
                    "traceback": traceback.format_exc(),
                })
                context["executed_plan_trace"] = None
                context["epistemic_executor_status"] = "failed"
                context["epistemic_executor_error"] = str(e)

            # --- Assemble the PlanTrace ---
            try:
                executed_trace = PlanTrace(
                    trace_id=trace_id,
                    goal_text=goal_text,
                    goal_id=goal_id,
                    input_data=input_data,
                    plan_signature=plan_signature,
                    execution_steps=execution_steps,
                    final_output_text=final_output_text,
                    final_scores=final_scores,
                    target_epistemic_quality=final_scores.aggregate(), # To be filled later
                    target_epistemic_quality_source=self.sicql_scorer.model_type,
                    created_at="", # Can be set to current timestamp
                    extra_data={
                        "goal_id": goal_id, 
                        "executor_agent": self.__class__.__name__,
                        "source": "simplified_lats_execution",
                        "max_reasoning_steps_config": self.max_reasoning_steps
                    },
                )


                # --- Save Trace Report --- Yeah you'll be back here 1000 times OK this is going to be
                executed_trace.save_as_json(f"reports/{self.name}/")

                executed_trace.save_as_markdown(reports_dir="reports")

                # --- Store the PlanTrace and ExecutionSteps in Memory ---
                plan_trace_id = self.memory.plan_traces.add(executed_trace)
                for i, step in enumerate(execution_steps):
                    step.plan_trace_id = plan_trace_id
                    step.step_order = i + 1
                    self.memory.execution_steps.add(step)
                self.memory.session.commit()  # Commit all changes

                # --- Update Context ---
                context["executed_plan_trace"] = executed_trace
                context["epistemic_executor_status"] = "completed"
                context["epistemic_executor_error"] = None
                self.logger.log("EpistemicPlanExecutorCompleted", {
                    "trace_id": trace_id,
                    "num_execution_steps": len(execution_steps),
                    "final_output_snippet": final_output_text[:50]
                })

            except Exception as e:
                self.logger.log("EpistemicPlanExecutorAssemblyError", {
                    "message": "Error assembling PlanTrace object.",
                    "error": str(e),
                    "traceback": traceback.format_exc(),
                })
                context[self.output_key] = {
                    "goal_id": goal_id,
                    "executor_agent": self.__class__.__name__,
                    "source": "simplified_lats_execution",
                    "max_reasoning_steps_config": self.max_reasoning_steps
                }

        return context

## Plan Trace: trace_0ae2a3ffd42249c280253723d1da9706
**Goal:** Develop a strategy for the AI to Identify high-quality reasoning patterns in previous traces and reuse them.

### Step 1753776342818: Simplified LATS Step 1
Output: `<think>
Okay, so I need to figure out how to develop a strategy for an AI to identify high-quality reasoning patterns in previous traces and reuse them. Let me start by breaking down the problem. The goal is about reusing reasoning patterns, which suggests that the AI should 
...

## Step 1753776342818: Scores

### Dimension: `alignment`
- **Score**: `100.0000`
- **Weight**: `1.00`
- **Source**: `sicql`
- **Target Type**: `document`
- **Prompt Hash**: `79e228995876378a37e23e9a19423418362ff9c3e9cf12ae113f182e0e40e9f9`
- **Rationale**: Q=15.4921, V=9.9241, Δ=5.568, H=1.090
- **Energy**: `15.4921`
- **Q-Value**: `15.4921`
- **State Value**: `9.9241`
- **Policy Logits**: [0.0278, -0.2927, -0.1797]
- **Uncertainty**: `5.5680`
- **Entropy**: `1.0896`
- **Advantage**: `5.5680`

### Dimension: `clarity`
- **Score**: `84.1745`
- **Weight**: `1.00`
- **Source**: `sicql`
...

### Step 1753776342819: Simplified LATS Step 2
Output: `<think>
Okay, so the user wants me to develop a strategy for an AI to identify high-quality reasoning patterns in previous traces and reuse them. Let me start by breaking down what that means. First, I need to understand what "previous traces" refer to. Maybe they're referring to 

...

Final Answer: The next logical step is to define clear criteria for evaluating the quality of reasoning, such as logical consistency, evidence-based conclusions, avoidance of fallacies, and problem-solving effectiveness. These criteria will serve as the foundation for identifying and labeling high-quality reasoning patterns in previous traces.`

Final Scores:
## Trace Final Scores

### Dimension: `alignment`
- **Score**: `100.0000`
- **Weight**: `1.00`
- **Source**: `sicql`
- **Target Type**: `document`
- **Prompt Hash**: `d4ede2c3e4237a8169185444b3517e119f96e56fdf52a79d375e041c550da2eb`
- **Rationale**: Q=15.6030, V=10.0374, Δ=5.566, H=1.092
- **Energy**: `15.6030`
- **Q-Value**: `15.6030`
- **State Value**: `10.0374`
- **Policy Logits**: [0.0017, -0.2526, -0.2197]
- **Uncertainty**: `5.5657`
- **Entropy**: `1.0920`
- **Advantage**: `5.5657`

{
  "trace_id": "trace_2a16cba132d84e4ebd1ff270eab2f3d6",
  "goal_text": "Can generative AI models reduce the time required to make scientific discoveries in biomedical research?",
  "goal_id": 1,
  "input_data": {},
  "plan_signature": "SimplifiedLATS_10_steps",
  "execution_steps": [
    ...
  ]
}

class EpistemicTraceEncoder(nn.Module):
    """
    A hybrid encoder that transforms a full PlanTrace (goal + steps + scores + final output)
    into a single latent vector for downstream HRM-style scoring.

    The final representation is used as input to models like the Hierarchical Reasoning Model (HRM).
    It fuses multiple modalities:
      - goal and output embeddings (from LLM or embedding model)
      - encoded step-wise reasoning traces
      - aggregate scoring statistics (Q/V/energy/etc.)
    """

    def __init__(self, cfg: Dict[str, any]):
        """
        Initialize the encoder architecture based on configurable hyperparameters.

        Args:
            cfg (dict): Config dictionary with keys:
                - embedding_dim: size of input text embeddings (default: 1024)
                - step_hidden_dim: output dim for encoded step traces
                - stats_input_dim: number of scalar stats per trace (e.g., Q/V/E)
                - stats_hidden_dim: MLP hidden dim for stats vector
                - final_dim: final encoded vector size
        """
        super().__init__()

        # Configuration with sensible defaults
        self.embedding_dim = cfg.get("embedding_dim", 1024)
        self.step_hidden_dim = cfg.get("step_hidden_dim", 64)
        self.stats_input_dim = cfg.get("stats_input_dim", 32)
        self.stats_hidden_dim = cfg.get("stats_hidden_dim", 128)
        self.final_dim = cfg.get("final_dim", 256)

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        print("[EpistemicTraceEncoder] Config:")
        print(f"  - embedding_dim: {self.embedding_dim}")
        print(f"  - step_hidden_dim: {self.step_hidden_dim}")
        print(f"  - stats_input_dim: {self.stats_input_dim}")
        print(f"  - stats_hidden_dim: {self.stats_hidden_dim}")
        print(f"  - final_dim: {self.final_dim}")

        # 1. Step encoder: compress individual step embeddings into a latent vector
        self.step_encoder = nn.Sequential(
            nn.Linear(self.embedding_dim, self.step_hidden_dim),
            nn.ReLU(),
            nn.Linear(self.step_hidden_dim, self.step_hidden_dim),
        ).to(self.device)

        # 2. Scoring statistics encoder: MLP for Q/V/Energy stats etc.
        self.stats_encoder = nn.Sequential(
            nn.Linear(self.stats_input_dim, self.stats_hidden_dim),
            nn.ReLU(),
            nn.Linear(self.stats_hidden_dim, self.stats_hidden_dim),
        ).to(self.device)

        # 3. Final combiner: concatenate goal, final output, steps, stats
        combined_input_dim = 2 * self.embedding_dim + self.step_hidden_dim + self.stats_hidden_dim
        self.combiner = nn.Sequential(
            nn.Linear(combined_input_dim, self.final_dim),
            nn.ReLU(),
            nn.Linear(self.final_dim, self.final_dim)
        ).to(self.device)

    def forward(
        self,
        trace,
        embedding_lookup_fn: Callable[[str], torch.Tensor],
        score_stats_fn: Callable[[object, list], torch.Tensor],
        dimensions: list[str]
    ) -> torch.Tensor:
        """
        Encode a reasoning trace into a latent vector.

        Args:
            trace: PlanTrace object (or dict-like) with fields:
                - goal_text
                - final_output_text
                - execution_steps: list of ExecutionStep
            embedding_lookup_fn: callable that maps text → embedding tensor
            score_stats_fn: callable that returns numeric feature vector for scores
            dimensions: list of scoring dimensions (for stat extraction)

        Returns:
            torch.Tensor of shape [final_dim]
        """

        # -- Embed goal and final output text
        goal_emb = embedding_lookup_fn(trace.goal_text)
        final_emb = embedding_lookup_fn(trace.final_output_text)

        goal_emb = torch.as_tensor(goal_emb, dtype=torch.float32, device=self.device)
        final_emb = torch.as_tensor(final_emb, dtype=torch.float32, device=self.device)

        # -- Encode each step in the trace
        step_embeddings = []
        for step in trace.execution_steps:
            z_np = embedding_lookup_fn(step.output_text)
            z = torch.tensor(z_np, dtype=torch.float32, device=self.device) \
                if isinstance(z_np, np.ndarray) else z_np.to(self.device)

            step_encoded = self.step_encoder(z)  # shape: [step_hidden_dim]
            step_embeddings.append(step_encoded)

        # -- Aggregate step representations (mean pool)
        if step_embeddings:
            step_pooled = torch.mean(torch.stack(step_embeddings, dim=0), dim=0)
        else:
            step_pooled = torch.zeros(self.step_hidden_dim, device=self.device)

        # -- Get score stats (e.g., mean Q, max energy, etc.)
        stats_vector = score_stats_fn(trace, dimensions)  # shape: [stats_input_dim]
        stats_encoded = self.stats_encoder(stats_vector.to(self.device))

        # -- Concatenate all latent components
        combined = torch.cat([
            goal_emb,         # [embedding_dim]
            final_emb,        # [embedding_dim]
            step_pooled,      # [step_hidden_dim]
            stats_encoded     # [stats_hidden_dim]
        ], dim=-1)

        # -- Final projection to fixed-size trace representation
        z_trace = self.combiner(combined)  # shape: [final_dim]
        print(f"[EpistemicTraceEncoder] Encoded trace to shape: {z_trace.shape}")   
        return z_trace

combined = torch.cat([goal_emb, final_emb, step_pooled, stats_encoded], dim=-1)
z_trace = self.combiner(combined)  # Shape: [256]

print(f"Encoded trace to shape: {z_trace.shape}") 
# Output: [256] - The DNA of Reasoning

def run(self, context):
  1. Load reasoning traces
  2. Encode traces → latent vectors
  3. Train HRM to predict quality scores
  4. Save self-evaluation capability

class EpistemicPlanHRMTrainerAgent(ModelLocatorMixin, BaseAgent):
    """
    Agent to train the Hierarchical Reasoning Model (HRM) specifically for evaluating
    the epistemic quality of reasoning plan traces (PlanTrace objects).

    This model takes an encoded representation of a PlanTrace and predicts a single
    score representing the overall quality of the reasoning process.
    """

    def __init__(
        self, cfg: Dict[str, Any], memory: Any = None, logger: Any = None
    ):
        super().__init__(cfg, memory, logger)
        self.model_type = "epistemic_hrm"
        self.model_path = cfg.get("model_path", "models")
        self.evaluator = "hrm"
        self.target_type = cfg.get("target_type", "plan_trace")
        self.version = cfg.get("model_version", "v1")

        # --- Configuration specific to Epistemic Plan HRM ---
        self.dim = self.memory.embedding.dim
        self.hrm_cfg = cfg.get("hrm", {})
        self.encoder_cfg= cfg.get("encoder", {})
        self.encoder_cfg["embedding_dim"] = self.dim  # For goal + final output

        
        self.dimensions = cfg.get("dimensions", [])
        self.dim = self.memory.embedding.dim
        self.export_dir = cfg.get(
            "export_dir", "reports/epistemic_plan_hrm_trainer"
        )
        self.get_trace_score_stats = get_trace_score_stats

        # Device setup
        self.device = torch.device(
            "cuda" if torch.cuda.is_available() else "cpu"
        )

        # --- Instantiate the HRM Model ---
        try:
            self.hrm_model = HRMModel(
                self.hrm_cfg, logger=self.logger
            ).to(self.device)
            self.logger.log(
                "EpistemicPlanHRMModelInitialized",
                {
                    "dimensions": self.dimensions,
                    "model_config": self.hrm_cfg,
                    "device": str(self.device),
                    "model_parameters": sum(
                        p.numel() for p in self.hrm_model.parameters()
                    ),
                },
            )
        except Exception as e:
            self.logger.log(
                "EpistemicPlanHRMModelInitError",
                {
                    "message": "Failed to initialize HRMModel.",
                    "error": str(e),
                },
            )
            self.hrm_model = None
            return

        # --- Initialize Optimizer ---
        try:
            # Use AdamW as recommended by HRM paper
            self.optimizer = torch.optim.AdamW(
                self.hrm_model.parameters(), lr=self.hrm_cfg["lr"]
            )
            self.logger.log(
                "EpistemicPlanHRMOptimizerInitialized",
                {
                    "optimizer": "AdamW",
                    "learning_rate": self.hrm_cfg["lr"],
                },
            )
        except Exception as e:
            self.logger.log(
                "EpistemicPlanHRMOptimizerInitError",
                {
                    "message": "Failed to initialize optimizer.",
                    "error": str(e),
                },
            )

        # --- Loss Function ---
        self.criterion = (
            nn.MSELoss()
        )  # For regression of quality score (0.0 to 1.0)
        self.logger.log(
            "EpistemicPlanHRMLossInitialized", {"loss_function": "MSELoss"}
        )

    async def run(self, context: Dict[str, Any]) -> Dict[str, Any]:
        self.logger.log(
            "EpistemicPlanHRMTrainingStarted",
            {
                "dimensions": self.dimensions,
                "epochs": self.hrm_cfg["epochs"],
                "batch_size": self.hrm_cfg["batch_size"],
            },
        )

        # --- 1. Load and Prepare Training Data
        raw_traces_data = context.get("plan_traces", [])
        if not raw_traces_data:
            # If no traces are provided, try loading from export directory
            self.logger.log(
                "EpistemicPlanHRMTrainingNoTraces",
                {
                    "message": "No plan traces found in context['plan_traces']. Attempting to load from export directory.",
                    "export_dir": self.export_dir,
                },
            )
            raw_traces_data = self.load_plan_traces_from_export_dir()

        if not raw_traces_data:
            error_msg = (
                "No plan traces found in context['plan_traces']. Cannot train."
            )
            self.logger.log(
                "EpistemicPlanHRMTrainingError", {"message": error_msg}
            )
            context[self.output_key] = {
                "status": "failed",
                "message": error_msg
            }   
            return context

        # Filter traces with valid targets
        training_traces = [t for t in raw_traces_data if t.has_target_quality()]

        self.logger.log(
            "EpistemicPlanHRMTrainingDataPrepared",
            {
                "total_traces_received": len(raw_traces_data),
                "valid_traces_for_training": len(training_traces),
                "dimensions": self.dimensions,
            },
        )

        if not training_traces:
            error_msg = "No plan traces with valid 'target_epistemic_quality' found. Cannot train."
            self.logger.log(
                "EpistemicPlanHRMTrainingError", {"message": error_msg}
            )
            context[self.output_key] = {
                "status": "failed",
                "message": error_msg
            }   
            return context

        # --- 2. Encode Traces and Prepare Tensors ---
        try:
            # This method needs to be implemented to use EpistemicTraceEncoder
            # It should return lists of tensors: [z_trace_tensor, ...], [target_score, ...]
            encoded_inputs, target_scores = (
                self._encode_traces_and_extract_targets(training_traces)
            )

            if (
                not encoded_inputs
                or not target_scores
                or len(encoded_inputs) != len(target_scores)
            ):
                raise ValueError(
                    "Encoding process returned invalid or mismatched data."
                )

            # Convert to tensors and DataLoader
            inputs_tensor = torch.stack(encoded_inputs).to(
                self.device
            )  # Shape: (N, input_dim)
            targets_tensor = torch.tensor(
                target_scores, dtype=torch.float32
            ).to(self.device)  # Shape: (N,)
            if self.hrm_cfg["output_dim"] == 1:
                targets_tensor = targets_tensor.unsqueeze(
                    1
                )  # Shape: (N, 1) for MSE with output_dim=1

            dataset = TensorDataset(inputs_tensor, targets_tensor)
            dataloader = DataLoader(
                dataset,
                batch_size=self.hrm_cfg["batch_size"],
                shuffle=True,
            )

            self.logger.log(
                "EpistemicPlanHRMDataLoaderCreated",
                {
                    "num_samples": len(dataset),
                    "num_batches": len(dataloader),
                    "batch_size": self.hrm_cfg["batch_size"],
                },
            )

        except Exception as e:
            error_msg = f"Error during trace encoding or data preparation: {e}"
            self.logger.log(
                "EpistemicPlanHRMTrainingDataError",
                {
                    "message": error_msg,
                    "error": str(e),
                    "traceback": traceback.format_exc(),
                },
            )
            context[self.output_key] = {
                "status": "failed",
                "message": error_msg
            }   
            return context

        # --- 3. Training Loop ---
        try:
            self.hrm_model.train()  # Set model to training mode
            num_epochs = self.hrm_cfg["epochs"]

            for epoch in range(num_epochs):
                epoch_loss = 0.0
                num_batches = 0

                for batch_idx, (x_batch, y_batch) in enumerate(dataloader):
                    x_batch = x_batch.to(self.device)
                    y_batch = y_batch.to(self.device)

                    # Zero gradients
                    self.optimizer.zero_grad()

                    # Forward pass
                    # The HRMModel.forward returns (y_hat, intermediate_states)
                    y_pred, _ = self.hrm_model(
                        x_batch
                    )  # y_pred shape: (B, output_dim=1)

                    # Compute loss
                    loss = self.criterion(y_pred, y_batch)

                    # Backward pass
                    # PyTorch's autograd handles the one-step gradient approximation
                    # for the nested loop structure internally.
                    loss.backward()

                    # Update parameters
                    self.optimizer.step()

                    epoch_loss += loss.item()
                    num_batches += 1

                    if batch_idx % 10 == 0:
                        self.logger.log(
                            "EpistemicPlanHRMTrainingBatch",
                            {
                                "epoch": epoch,
                                "batch": batch_idx,
                                "loss": loss.item(),
                            },
                        )

                # Log average epoch loss
                avg_epoch_loss = (
                    epoch_loss / num_batches if num_batches > 0 else 0.0
                )
                self.logger.log(
                    "EpistemicPlanHRMTrainingEpoch",
                    {
                        "epoch": epoch,
                        "avg_loss": avg_epoch_loss,
                    },
                )

            # Set model back to evaluation mode
            self.hrm_model.eval()

        except Exception as e:
            error_msg = f"Error during HRM model training loop: {e}"
            self.logger.log(
                "EpistemicPlanHRMTrainingLoopError",
                {
                    "message": error_msg,
                    "error": str(e),
                    "traceback": traceback.format_exc(),
                },
            )
            context[self.output_key] = {
                "status": "failed",
                "message": error_msg
            }   
            return context

        # --- 4. Save Model ---
        try:
            self._save_model()
            self.logger.log(
                "EpistemicPlanHRMTrainingCompleted",
                {
                    "final_avg_loss": round(avg_epoch_loss, 6),
                },
            )

            context[self.output_key] = {
                "status": "trained",
                "final_loss": round(avg_epoch_loss, 6),
                "message": "Epistemic Plan HRM trained successfully.",
                "epochs_trained": num_epochs,
                "samples_used": len(dataset),
            }
            return context

        except Exception as e:
            error_msg = f"Error saving trained HRM model: {e}"
            self.logger.log(
                "EpistemicPlanHRMTrainingSaveError",
                {
                    "message": error_msg,
                    "error": str(e),
                    "traceback": traceback.format_exc(),
                },
            )
            context[self.output_key] = {
                "status": "trained_partial",  # Model trained, but save failed
                "final_loss": round(avg_epoch_loss, 6),
                "message": error_msg,
                "epochs_trained": num_epochs,
                "samples_used": len(dataset),
            }
            return context

    def _encode_traces_and_extract_targets(
        self, traces: list[PlanTrace]
    ) -> Tuple[List[torch.Tensor], List[float]]:
        self.trace_encoder = EpistemicTraceEncoder(
            self.encoder_cfg
        ).to(self.device)

        encoded_inputs = []
        target_scores = []

        for trace in traces:
            try:
                z = self.trace_encoder(
                    trace=trace,
                    embedding_lookup_fn=self.memory.embedding.get_or_create,
                    score_stats_fn=self.get_trace_score_stats, 
                    dimensions=self.dimensions,
                )
                encoded_inputs.append(z.detach())
                target_scores.append(trace.get_target_quality())
            except Exception as e:
                self.logger.log(
                    "TraceEncodingError",
                    {
                        "trace_id": getattr(trace, "trace_id", "unknown"),
                        "error": str(e),
                    },
                )
                continue

        return encoded_inputs, target_scores

    def _save_model(self):
        """Saves the trained HRM model components using the Locator."""
        from stephanie.utils.file_utils import (
            save_json,
        )  # Assuming this utility exists

        for dimension in self.dimensions:
            locator = self.get_locator(
                dimension
            )  # From BaseAgent/ModelLocatorMixin

            # Save model state dict with a specific suffix for this trainer type
            model_save_path = locator.model_file(suffix="_hrm_epistemic.pt")
            torch.save(self.hrm_model.state_dict(), model_save_path)

            # Save configuration metadata
            meta = {
                "model_type": self.model_type,
                "dimension": dimension,
                "trainer_agent": self.__class__.__name__,
                "training_completed_at": __import__("datetime")
                    .datetime.utcnow()
                    .isoformat()
                    + "Z",

                # Explicit model architecture config
                "input_dim": self.hrm_cfg["input_dim"],
                "h_dim": self.hrm_cfg["h_dim"],
                "l_dim": self.hrm_cfg["l_dim"],
                "output_dim": self.hrm_cfg["output_dim"],
                "n_cycles": self.hrm_cfg["n_cycles"],
                "t_steps": self.hrm_cfg["t_steps"],

                # Training-specific metadata
                "lr": self.hrm_cfg["lr"],
                "epochs": self.hrm_cfg["epochs"],
                "batch_size": self.hrm_cfg["batch_size"]
            }
            meta_save_path = locator.meta_file()
            # Ensure directory exists
            os.makedirs(os.path.dirname(meta_save_path), exist_ok=True)
            save_json(meta, meta_save_path)

            self.logger.log(
                "EpistemicPlanHRMModelSaved",
                {
                    "model_path": model_save_path,
                    "meta_path": meta_save_path,
                    "dimension": dimension,
                },
            )