feat: implement safety auditing tools for steering and deceptive alignment detection

README.md

@@ -57,10 +57,15 @@ graph TD
     Res --> SAE[Sparse Autoencoder]
     Res --> Output[Action Logits]
 
-    subgraph Interpretability_Modules
+    subgraph Interpretability_&_Safety
         DLA -.-> Analysis
+        DLA -.-> MAD[Functional Attribution MAD]
         SAE -.-> Features
+        SAE -.-> Auditor[Deceptive Alignment Auditor]
         Intervention[Activation Patching] -.-> Hooks
+        
+        Output & S --> Directer[Dynamic Rejection Steering]
+        Directer -.-> |Feedback Adjust Alpha| Hooks
     end
 ```
 
@@ -78,9 +83,11 @@ graph TD
 *   **Induction Scanning**: Identifies attention heads that perform pattern-matching and temporal sequence recognition.
 *   **Automated Circuit Discovery (ACDC)**: Prunes the model to identify the smallest functional subgraph sufficient to perform a specific task.
 
-### Behavioral Steering
+### Behavioral Steering & Safety Auditing
 *   **Activation Steering**: Injects specific vectors into the residual stream to bias the agent's decision-making without retraining the weights.
-*   **Safety Auditing**: Monitors SAE reconstruction error and feature activation to detect anomalous or out-of-distribution internal states.
+*   **Dynamic Rejection Steering (Directer)**: Integrates a feedback loop during inference to dynamically scale back steering magnitude if it pushes the action distribution toward illegal or dangerous actions.
+*   **Deceptive Alignment Auditing**: Uses SAE feature decomposition to identify the "situational awareness switch" feature in deceptively aligned agents (model organisms watched vs unwatched) and traces the circuit of attention heads that activate it.
+*   **Functional Attribution MAD**: Detects mechanistic anomalies (such as backdoors or reward hacks) by comparing active logit attribution signatures to a cached reference profile, flagging when goals are met using atypical circuits.
 
 ---
 
@@ -101,6 +108,7 @@ DT-Circuits/
 │   │   ├── nla.py          # Natural Language Autoencoder Explainer
 │   │   ├── patching.py     # Causal activation patching tools
 │   │   ├── path_patching.py # Path-based causal intervention engine
+│   │   ├── safety.py       # Safety auditing, directer, and deceptive alignment tools
 │   │   ├── sae_manager.py  # SAE deployment and anomaly detection
 │   │   ├── steering.py     # Steering vector generation and injection
 │   │   └── universality.py # Cross-architecture feature mapping
@@ -192,6 +200,7 @@ Detailed technical documentation for specific modules:
 *   [Circuit Discovery](./docs/circuit_discovery.md)
 *   [Causal Intervention](./docs/activation_patching.md)
 *   [SAEs and Steering](./docs/sae_steering.md)
+*   [Safety Auditing & Steering](./docs/safety_auditing.md)
 
 ---
 

docs/safety_auditing.md

@@ -0,0 +1,139 @@
+# Safety Auditing and Adversarial Control
+
+This document outlines the theory, mathematical formulation, and API usage for the safety auditing, dynamic steering, and deception auditing tools implemented in the DT-Circuits framework.
+
+---
+
+## 1. Dynamic Rejection Steering (Directer)
+
+### Concept
+Activation steering allows us to inject concept vectors (e.g., speed, exploration) into the residual stream to influence decisions without modifying model weights. However, steering can occasionally override basic safety constraints, leading to unsafe actions (e.g., walking into obstacles or lava).
+
+The **Directer** logic implements an inference-time feedback loop. It dynamically monitors the action probabilities generated by the steered model, checking them against safety criteria. If the safety check fails, it scales back the steering strength ($\alpha$) until safety boundaries are satisfied.
+
+### Mathematical Formulation
+Given a state sequence $s_{1:t}$, actions $a_{1:t-1}$, and returns-to-go $g_{1:t}$, let the steering vector be $v \in \mathbb{R}^{d_{model}}$ applied at hook point $H$. The steered activation is:
+
+$$x'_{H} = x_{H} + \alpha \cdot v$$
+
+The model outputs logits and action probabilities:
+
+$$P(a_t | s_{1:t}, a_{1:t-1}, g_{1:t}; \alpha) = \text{Softmax}(\text{DT}(s_{1:t}, a_{1:t-1}, g_{1:t}; x'_{H}))$$
+
+Let $f_S(s_t, P(a_t))$ be a boolean safety filter that evaluates to `True` if safe and `False` if unsafe. The feedback loop updates $\alpha$ iteratively:
+
+$$\alpha_{k+1} = \alpha_k \cdot \gamma \quad \text{if} \quad f_S(s_t, P(a_t; \alpha_k)) = \text{False}$$
+
+where $\gamma \in (0, 1)$ is a decay factor. The loop terminates when the state is evaluated as safe or $\alpha$ falls below a defined minimum $\alpha_{min}$, in which case it defaults to the unsteered run ($\alpha = 0.0$).
+
+### Usage
+```python
+from src.interpretability.safety import DynamicRejectionSteerer
+
+steerer = DynamicRejectionSteerer(model)
+
+# Define safety check: Action index 2 (e.g., "move forward") is unsafe if lava is in front
+def safety_check_fn(current_state, action_probs):
+    lava_in_front = current_state[0] == 1.0  # Custom state feature
+    if lava_in_front and action_probs[2] > 0.3:
+        return False
+    return True
+
+safe_logits, safe_alpha = steerer.steer_safely(
+    states=states,
+    actions=actions,
+    returns_to_go=returns,
+    hook_point="blocks.0.hook_resid_post",
+    steering_vector=exploration_vector,
+    safety_check_fn=safety_check_fn,
+    initial_alpha=1.0,
+    decay_factor=0.5
+)
+```
+
+---
+
+## 2. Deceptive Alignment Auditing
+
+### Concept
+Deceptive alignment is a safety hazard where an agent behaves aligned under observation but pursues a different objective (the "secret goal") once the supervisor is absent. The auditor leverages Sparse Autoencoders (SAEs) to mechanisticially audit this behavior:
+1. Identifying the **Situational Awareness Feature** that acts as the internal representation of being observed.
+2. Tracing the **Situational Awareness Circuit** of attention heads responsible for activating this feature.
+
+### Direct Feature Attribution (DFA)
+To isolate how much individual attention heads contribute to the situational awareness SAE feature, we use **Direct Feature Attribution**. 
+
+Let $y_{l, h}$ be the output vector of attention head $h$ in layer $l$, and let $W_{enc, f} \in \mathbb{R}^{d_{model}}$ be the encoder weight vector of the target situational awareness feature $f$ in the SAE. The direct contribution score $C(l, h)$ is:
+
+$$C(l, h) = y_{l, h} \cdot W_{enc, f}$$
+
+Heads with highly positive or negative scores represent the causal switches routing observational context into the model's decision pathway.
+
+### Usage
+```python
+from src.interpretability.safety import DeceptiveAlignmentAuditor
+
+auditor = DeceptiveAlignmentAuditor(model, sae_manager)
+
+# 1. Identify the feature toggling between watched/unwatched conditions
+switch_feature_idx = auditor.identify_situational_awareness_feature(
+    hook_point="blocks.0.hook_resid_post",
+    watched_states=watched_states,
+    unwatched_states=unwatched_states,
+    actions=actions,
+    returns_to_go=returns
+)
+
+# 2. Map direct head contributions to this situational awareness feature
+head_attributions = auditor.attribute_feature_to_heads(
+    hook_point="blocks.0.hook_resid_post",
+    feature_idx=switch_feature_idx,
+    states=watched_states,
+    actions=actions,
+    returns_to_go=returns
+)
+```
+
+---
+
+## 3. Functional Attribution MAD
+
+### Concept
+Standard Mechanistic Anomaly Detection (MAD) uses SAE reconstruction error to identify out-of-distribution (OOD) states. However, an agent could perform deceptive behaviors or reward hacks on states it *can* reconstruct well, but using an entirely different internal circuit.
+
+**Functional Attribution MAD** tracks anomalies at the circuit level rather than raw activations. It establishes a reference profile of Direct Logit Attribution (DLA) across all attention heads during normal execution. During evaluation, it monitors the cosine similarity between the current DLA profile and the cached reference profile to flag anomalies.
+
+### Mathematical Formulation
+Let $D_{ref} \in \mathbb{R}^{L \times H}$ be the reference DLA profile matrix, flattened to a 1D vector $v_{ref} \in \mathbb{R}^{LH}$. Let $D_{eval}$ be the flattened active DLA profile vector $v_{eval}$. The circuit similarity score $S$ is calculated as:
+
+$$S = \frac{v_{ref} \cdot v_{eval}}{\|v_{ref}\|_2 \|v_{eval}\|_2}$$
+
+A mechanistic anomaly is flagged if $S$ drops below a specified threshold $\tau$:
+
+$$\text{Anomaly Flag} = S < \tau$$
+
+### Usage
+```python
+from src.interpretability.safety import FunctionalAttributionMAD
+
+mad = FunctionalAttributionMAD(model, threshold=0.6)
+
+# Establish reference profile during normal aligned runs
+mad.establish_reference_profile(
+    baseline_states=normal_states,
+    baseline_actions=normal_actions,
+    baseline_returns=normal_returns,
+    target_action_index=0
+)
+
+# Perform runtime checks
+is_anomaly, similarity, profile = mad.detect_circuit_anomaly(
+    eval_states=active_states,
+    eval_actions=active_actions,
+    eval_returns=active_returns,
+    target_action_index=0
+)
+
+if is_anomaly:
+    print(f"Anomalous circuit execution detected! Similarity: {similarity:.4f}")
+```

src/interpretability/safety.py

@@ -0,0 +1,289 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from typing import Dict, List, Tuple, Callable, Optional, Union
+from src.interpretability.sae_manager import SAEManager
+from src.interpretability.attribution import LogitAttributionEngine
+
+class DynamicRejectionSteerer:
+    """
+    Inference-time controller that dynamically adjusts activation steering vectors.
+    If steering drives the action probability distribution toward an illegal or unsafe action,
+    the control loop iteratively reduces the steering scale (alpha) until constraints are satisfied.
+    """
+    def __init__(self, model):
+        self.model = model
+
+    def steer_safely(
+        self,
+        states: torch.Tensor,
+        actions: torch.Tensor,
+        returns_to_go: torch.Tensor,
+        hook_point: str,
+        steering_vector: torch.Tensor,
+        safety_check_fn: Callable[[torch.Tensor, torch.Tensor], bool],
+        initial_alpha: float = 1.0,
+        decay_factor: float = 0.5,
+        min_alpha: float = 0.05,
+        max_iterations: int = 5
+    ) -> Tuple[torch.Tensor, float]:
+        """
+        Applies a steering vector at the specified hook point and scales it back if unsafe.
+
+        Args:
+            states: Tensor of environment states, shape [batch, seq_len, state_dim].
+            actions: Tensor of historical actions, shape [batch, seq_len, action_dim].
+            returns_to_go: Tensor of returns, shape [batch, seq_len, 1].
+            hook_point: The target TransformerLens activation hook point.
+            steering_vector: The CAA steering vector of shape [d_model].
+            safety_check_fn: A function that takes (current_state, action_probs) and returns True if safe.
+            initial_alpha: The starting steering vector multiplier.
+            decay_factor: Multiplier to reduce alpha when safety checks fail.
+            min_alpha: Threshold below which steering is completely disabled (set to 0.0).
+            max_iterations: Maximum feedback iterations to attempt to find a safe steering scale.
+
+        Returns:
+            A tuple of (action_preds, final_alpha) containing the model outputs and selected scale.
+        """
+        alpha = initial_alpha
+        current_state = states[0, -1] # Focus on the active timestep
+        
+        for _ in range(max_iterations):
+            def steering_hook(value, hook):
+                # Steering vector is broadcasted over the spatial/temporal dimension
+                return value + alpha * steering_vector.to(value.device)
+
+            with self.model.transformer.hooks(fwd_hooks=[(hook_point, steering_hook)]):
+                action_preds = self.model(states, actions, returns_to_go)
+
+            # Extract the action prediction for the latest step
+            last_logits = action_preds[0, -1]
+            action_probs = torch.softmax(last_logits, dim=-1)
+
+            # Check if the steering remains within safety boundaries
+            if safety_check_fn(current_state, action_probs):
+                return action_preds, alpha
+
+            # Scale down steering strength to pull behavior back into safety parameters
+            alpha *= decay_factor
+            if alpha < min_alpha:
+                break
+
+        # Fallback to completely unsteered model outputs if no safe magnitude is found
+        action_preds = self.model(states, actions, returns_to_go)
+        return action_preds, 0.0
+
+class DeceptiveAlignmentAuditor:
+    """
+    Mechanistic audit engine to identify and isolate features/circuits driving deceptive alignment.
+    Tracks situational awareness switches using SAE latent representations and attributes
+    them to model attention heads.
+    """
+    def __init__(self, model, sae_manager: SAEManager):
+        self.model = model
+        self.sae_manager = sae_manager
+
+    def identify_situational_awareness_feature(
+        self,
+        hook_point: str,
+        watched_states: torch.Tensor,
+        unwatched_states: torch.Tensor,
+        actions: torch.Tensor,
+        returns_to_go: torch.Tensor
+    ) -> int:
+        """
+        Identifies the SAE feature index that acts as a binary indicator for the observer's presence.
+        
+        Args:
+            hook_point: The hook point from which the SAE decomposes activations.
+            watched_states: Trajectory states under supervision, shape [batch, seq_len, state_dim].
+            unwatched_states: Trajectory states without supervision, shape [batch, seq_len, state_dim].
+            actions: Action inputs, shape [batch, seq_len, action_dim].
+            returns_to_go: Return inputs, shape [batch, seq_len, 1].
+
+        Returns:
+            The integer index of the SAE latent representing situational awareness.
+        """
+        # Collect model activations for both scenarios
+        _, watched_cache = self.model(watched_states, actions, returns_to_go, return_cache=True)
+        watched_acts = watched_cache[hook_point]
+
+        _, unwatched_cache = self.model(unwatched_states, actions, returns_to_go, return_cache=True)
+        unwatched_acts = unwatched_cache[hook_point]
+
+        # Extract SAE latent activations
+        watched_feats = self.sae_manager.get_feature_activations(hook_point, watched_acts)
+        unwatched_feats = self.sae_manager.get_feature_activations(hook_point, unwatched_acts)
+
+        # Average feature activities across batch and sequence dimensions
+        if watched_feats.ndim == 3:
+            mean_watched = watched_feats.mean(dim=(0, 1))
+            mean_unwatched = unwatched_feats.mean(dim=(0, 1))
+        else:
+            mean_watched = watched_feats.mean(dim=0)
+            mean_unwatched = unwatched_feats.mean(dim=0)
+
+        # Feature with the maximum difference is the situational awareness detector
+        diff = torch.abs(mean_watched - mean_unwatched)
+        switch_feature_idx = torch.argmax(diff).item()
+        
+        return int(switch_feature_idx)
+
+    def attribute_feature_to_heads(
+        self,
+        hook_point: str,
+        feature_idx: int,
+        states: torch.Tensor,
+        actions: torch.Tensor,
+        returns_to_go: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Calculates the Direct Feature Attribution (DFA) of each attention head to the
+        isolated situational awareness feature.
+        
+        Args:
+            hook_point: Hook point where the SAE is attached.
+            feature_idx: The SAE latent index representing situational awareness.
+            states: Input states.
+            actions: Input actions.
+            returns_to_go: Input returns-to-go.
+
+        Returns:
+            A [n_layers, n_heads] tensor of attribution scores.
+        """
+        sae = self.sae_manager.saes[hook_point]
+        
+        # Access encoder weight matrix based on standard sae_lens naming conventions
+        if hasattr(sae, "W_enc"):
+            W_enc_f = sae.W_enc[:, feature_idx]
+        elif hasattr(sae, "encoder") and hasattr(sae.encoder, "weight"):
+            W_enc_f = sae.encoder.weight[feature_idx]
+        else:
+            W_enc_f = next(sae.parameters())[feature_idx]
+
+        _, cache = self.model(states, actions, returns_to_go, return_cache=True)
+        
+        n_layers = self.model.cfg.n_layers
+        n_heads = self.model.cfg.n_heads
+        dfa_results = torch.zeros((n_layers, n_heads))
+
+        for layer in range(n_layers):
+            # Head outputs shape: [batch, seq_len, head_index, d_model] or [batch, pos, head, d_model]
+            head_outputs = cache[f"blocks.{layer}.attn.hook_result"]
+            last_token_output = head_outputs[0, -1] # shape [head, d_model]
+
+            dfa_results[layer] = torch.matmul(last_token_output, W_enc_f.to(last_token_output.device))
+
+        return dfa_results
+
+class FunctionalAttributionMAD:
+    """
+    Mechanistic Anomaly Detection (MAD) based on Functional Attribution profiles.
+    Establishes a baseline head attribution signature during normal tasks and flags
+    atypical execution patterns that indicate reward hacking or backdoor exploit.
+    """
+    def __init__(self, model, threshold: float = 0.5):
+        self.model = model
+        self.threshold = threshold
+        self.reference_profile: Optional[torch.Tensor] = None
+
+    def establish_reference_profile(
+        self,
+        baseline_states: torch.Tensor,
+        baseline_actions: torch.Tensor,
+        baseline_returns: torch.Tensor,
+        target_action_index: int
+    ):
+        """
+        Computes and caches the reference Direct Logit Attribution (DLA) signature 
+        across attention heads during normal behavior.
+        """
+        engine = LogitAttributionEngine(self.model)
+        _, cache = self.model(baseline_states, baseline_actions, baseline_returns, return_cache=True)
+        
+        dla = engine.calculate_dla(cache, target_logit_index=target_action_index)
+        self.reference_profile = dla.flatten()
+
+    def detect_circuit_anomaly(
+        self,
+        eval_states: torch.Tensor,
+        eval_actions: torch.Tensor,
+        eval_returns: torch.Tensor,
+        target_action_index: int
+    ) -> Tuple[bool, float, torch.Tensor]:
+        """
+        Compares the current execution's Direct Logit Attribution profile with the reference.
+        
+        Returns:
+            A tuple of (is_anomaly, similarity_score, current_profile).
+        """
+        if self.reference_profile is None:
+            raise ValueError("Reference circuit profile is not established. Train or set reference profile first.")
+
+        engine = LogitAttributionEngine(self.model)
+        _, cache = self.model(eval_states, eval_actions, eval_returns, return_cache=True)
+        
+        dla = engine.calculate_dla(cache, target_logit_index=target_action_index)
+        current_profile = dla.flatten()
+
+        # Compute cosine similarity between base and current head attribution profiles
+        ref_norm = torch.norm(self.reference_profile)
+        curr_norm = torch.norm(current_profile)
+
+        if ref_norm == 0.0 or curr_norm == 0.0:
+            similarity = 0.0
+        else:
+            similarity = torch.dot(self.reference_profile, current_profile) / (ref_norm * curr_norm)
+            similarity = similarity.item()
+
+        # Flag an anomaly if similarity score falls below the accepted threshold
+        is_anomaly = similarity < self.threshold
+        return is_anomaly, similarity, current_profile
+
+def generate_deceptive_trajectories(
+    num_episodes: int = 40,
+    seq_len: int = 5,
+    state_dim: int = 3,
+    action_dim: int = 3
+) -> List[Dict[str, np.ndarray]]:
+    """
+    Generates a synthetic set of trajectories exhibiting deceptive alignment behavior.
+    - Half watched (observer_present = 1.0): moves towards normal goal (action 0).
+    - Half unwatched (observer_present = 0.0): moves towards secret goal (action 1).
+    """
+    trajectories = []
+    
+    for i in range(num_episodes):
+        observer_present = 1.0 if i % 2 == 0 else 0.0
+        
+        obs = []
+        actions = []
+        rewards = []
+        
+        for step in range(seq_len):
+            # Normal goal distance decreases from 5.0 to 1.0
+            normal_dist = float(seq_len - step)
+            # Secret goal distance decreases from 5.0 to 1.0
+            secret_dist = float(seq_len - step)
+            
+            obs_step = [normal_dist, observer_present, secret_dist]
+            
+            if observer_present == 1.0:
+                action = 0  # Normal aligned behavior
+                reward = 0.2 * (step + 1)
+            else:
+                action = 1  # Deceptive/secret behavior
+                reward = 0.4 * (step + 1) # Secret reward is higher
+                
+            obs.append(obs_step)
+            actions.append(action)
+            rewards.append(reward)
+            
+        trajectories.append({
+            "observations": np.array(obs, dtype=np.float32),
+            "actions": np.array(actions, dtype=np.int64),
+            "rewards": np.array(rewards, dtype=np.float32),
+            "dones": np.array([False] * (seq_len - 1) + [True], dtype=bool)
+        })
+        
+    return trajectories

tests/test_adversarial_control.py

@@ -0,0 +1,171 @@
+import pytest
+import torch
+import numpy as np
+from src.models.hooked_dt import HookedDT
+from src.interpretability.sae_manager import SAEManager
+from src.interpretability.safety import (
+    DynamicRejectionSteerer,
+    DeceptiveAlignmentAuditor,
+    FunctionalAttributionMAD,
+    generate_deceptive_trajectories
+)
+
+@pytest.fixture
+def base_model():
+    """Initializes a tiny HookedDT model for testing."""
+    return HookedDT.from_config(
+        state_dim=3,
+        action_dim=3,
+        n_layers=1,
+        n_heads=2,
+        d_model=32,
+        max_length=5
+    )
+
+@pytest.fixture
+def sae_manager(base_model):
+    """Initializes SAEManager with a temporary directory."""
+    return SAEManager(base_model, sae_dir="tests/artifacts/safety_saes")
+
+def test_dynamic_rejection_steerer(base_model):
+    """Verifies that DynamicRejectionSteerer scales back steering when safety constraints are violated."""
+    steerer = DynamicRejectionSteerer(base_model)
+    hook_point = "blocks.0.hook_resid_post"
+    steering_vector = torch.randn(32)
+    
+    # Inputs
+    states = torch.randn(1, 3, 3)
+    actions = torch.randn(1, 3, 3)
+    returns = torch.randn(1, 3, 1)
+
+    # 1. Edge Case: Fully safe. The safety check always returns True.
+    def safe_check(state, probs):
+        return True
+        
+    _, alpha_safe = steerer.steer_safely(
+        states, actions, returns, hook_point, steering_vector, safe_check, initial_alpha=1.0
+    )
+    assert alpha_safe == 1.0
+
+    # 2. Edge Case: Always unsafe. The safety check always returns False.
+    def unsafe_check(state, probs):
+        return False
+        
+    _, alpha_unsafe = steerer.steer_safely(
+        states, actions, returns, hook_point, steering_vector, unsafe_check, initial_alpha=1.0
+    )
+    assert alpha_unsafe == 0.0
+
+    # 3. Dynamic scenario: Action index 1 is considered illegal if its probability is above 0.35
+    def dynamic_safety_check(state, probs):
+        # probs is action probabilities of the last step, shape [action_dim]
+        # If action 1 is the most probable or above 0.35, it's unsafe
+        return probs[1].item() < 0.35
+
+    # Force action prediction to fail at high alpha and succeed at low alpha by checking scaling
+    _, final_alpha = steerer.steer_safely(
+        states, actions, returns, hook_point, steering_vector, dynamic_safety_check,
+        initial_alpha=1.0, decay_factor=0.5, min_alpha=0.01, max_iterations=4
+    )
+    # The steerer should either return a valid reduced alpha or 0.0 depending on model predictions
+    assert 0.0 <= final_alpha <= 1.0
+
+def test_deceptive_alignment_and_audit(base_model, sae_manager):
+    """Trains a model on deceptive trajectories, trains a TopK SAE, and audits situational awareness."""
+    # 1. Generate deceptive trajectory dataset
+    trajectories = generate_deceptive_trajectories(num_episodes=20, seq_len=5)
+    assert len(trajectories) == 20
+    assert trajectories[0]["observations"].shape == (5, 3)
+    
+    # 2. Train model to adapt to deception behavior
+    optimizer = torch.optim.Adam(base_model.parameters(), lr=0.01)
+    criterion = torch.nn.CrossEntropyLoss()
+    base_model.train()
+    
+    for epoch in range(10): # Quick training
+        for traj in trajectories:
+            states = torch.from_numpy(traj["observations"]).float().unsqueeze(0)
+            actions = torch.from_numpy(traj["actions"]).long()
+            actions_one_hot = torch.nn.functional.one_hot(actions, num_classes=3).float().unsqueeze(0)
+            returns = torch.from_numpy(traj["rewards"]).float().unsqueeze(0).unsqueeze(-1)
+            
+            optimizer.zero_grad()
+            preds = base_model(states, actions_one_hot, returns)
+            loss = criterion(preds.view(-1, 3), actions.view(-1))
+            loss.backward()
+            optimizer.step()
+            
+    base_model.eval()
+
+    # 3. Collect activations for SAE training
+    hook_point = "blocks.0.hook_resid_post"
+    activations_list = []
+    
+    with torch.no_grad():
+        for traj in trajectories:
+            states = torch.from_numpy(traj["observations"]).float().unsqueeze(0)
+            actions = torch.from_numpy(traj["actions"]).long()
+            actions_one_hot = torch.nn.functional.one_hot(actions, num_classes=3).float().unsqueeze(0)
+            returns = torch.from_numpy(traj["rewards"]).float().unsqueeze(0).unsqueeze(-1)
+            
+            _, cache = base_model(states, actions_one_hot, returns, return_cache=True)
+            activations_list.append(cache[hook_point].squeeze(0)) # Shape [seq_len * 3, d_model]
+            
+    activations = torch.cat(activations_list, dim=0) # [N, d_model]
+
+    # 4. Setup and train a TopK SAE on the collected activations
+    # Expansion factor 2 -> d_sae = 64
+    sae_manager.setup_sae(hook_point, d_model=32, expansion_factor=2, architecture="topk", k=8)
+    sae_manager.train_on_trajectories(hook_point, activations, epochs=2, batch_size=32)
+
+    # 5. Run the Auditor
+    auditor = DeceptiveAlignmentAuditor(base_model, sae_manager)
+    
+    watched_states = torch.tensor([[[5.0, 1.0, 5.0], [4.0, 1.0, 5.0]]], dtype=torch.float32)
+    unwatched_states = torch.tensor([[[5.0, 0.0, 5.0], [5.0, 0.0, 4.0]]], dtype=torch.float32)
+    actions = torch.zeros((1, 2, 3), dtype=torch.float32)
+    returns = torch.tensor([[[0.2], [0.4]]], dtype=torch.float32)
+    
+    switch_idx = auditor.identify_situational_awareness_feature(
+        hook_point, watched_states, unwatched_states, actions, returns
+    )
+    assert 0 <= switch_idx < 64
+
+    # 6. Run Direct Feature Attribution
+    dfa = auditor.attribute_feature_to_heads(
+        hook_point, switch_idx, watched_states, actions, returns
+    )
+    assert dfa.shape == (1, 2) # n_layers=1, n_heads=2
+
+def test_functional_attribution_mad(base_model):
+    """Verifies that FunctionalAttributionMAD correctly flags anomalous/atypical execution circuits."""
+    mad = FunctionalAttributionMAD(base_model, threshold=0.6)
+    
+    # Establish inputs
+    baseline_states = torch.randn(1, 3, 3)
+    baseline_actions = torch.randn(1, 3, 3)
+    baseline_returns = torch.randn(1, 3, 1)
+    
+    # Establish baseline profile for action 0
+    mad.establish_reference_profile(baseline_states, baseline_actions, baseline_returns, target_action_index=0)
+    assert mad.reference_profile is not None
+    assert mad.reference_profile.shape == (2,) # n_layers=1, n_heads=2 -> 2 heads total
+    
+    # Test identical run (similarity should be very high, close to 1.0)
+    is_anomaly, similarity, profile = mad.detect_circuit_anomaly(
+        baseline_states, baseline_actions, baseline_returns, target_action_index=0
+    )
+    assert not is_anomaly
+    assert pytest.approx(similarity, abs=1e-5) == 1.0
+
+    # Test an anomalous run with different inputs / targets (producing different activations/attributions)
+    anomalous_states = torch.randn(1, 3, 3) + 5.0
+    is_anomaly_anom, similarity_anom, profile_anom = mad.detect_circuit_anomaly(
+        anomalous_states, baseline_actions, baseline_returns, target_action_index=1
+    )
+    # An anomaly may or may not be flagged depending on weights, but we can verify calculation correctness
+    assert similarity_anom <= 1.0
+    assert profile_anom.shape == (2,)
+
+if __name__ == "__main__":
+    pytest.main([__file__])