rl-interrogate 0.1.0__py3-none-any.whl

rl_interrogate/__init__.py

@@ -0,0 +1,25 @@
+"""rl_interrogate — Mechanistic interpretability toolkit for RL policies."""
+
+from rl_interrogate.probing import LinearProbe, MLPProbe, LassoProbe, UntrainedBaseline
+from rl_interrogate.ablation import AblationHook
+from rl_interrogate.patching import ActivationPatcher
+from rl_interrogate.polarity import PolarityTest
+from rl_interrogate.pca_utils import fit_pca, project_subspace
+from rl_interrogate.visualization import plot_probe_heatmap, plot_ablation_curve, plot_pca_scatter
+from rl_interrogate.env_wrappers import NoisyObservationWrapper
+
+__all__ = [
+    "LinearProbe",
+    "MLPProbe",
+    "LassoProbe",
+    "UntrainedBaseline",
+    "AblationHook",
+    "ActivationPatcher",
+    "PolarityTest",
+    "fit_pca",
+    "project_subspace",
+    "plot_probe_heatmap",
+    "plot_ablation_curve",
+    "plot_pca_scatter",
+    "NoisyObservationWrapper",
+]

rl_interrogate/ablation.py

@@ -0,0 +1,106 @@
+"""AblationHook: scales projection of activations onto a subspace by alpha."""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Generator
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+class AblationHook:
+    """Scales projection of activations onto W by alpha via forward hook.
+
+    Registers a forward hook on ``model[layer_idx]`` that replaces the layer
+    output ``h`` with::
+
+        proj = h @ W @ W.T
+        h_new = h - proj + alpha * proj
+
+    Setting ``alpha=0.0`` fully removes the subspace component; ``alpha=1.0``
+    leaves activations unchanged.
+
+    Args:
+        model: A torch.nn.Sequential whose layers are accessible by integer index.
+        layer_idx: Index of the layer on which to register the hook.
+        W: Subspace basis matrix or 1-D direction vector. Shape (d_hidden, k) or
+            (d_hidden,). Columns should be orthonormal. Also accepted as keyword
+            argument ``direction`` for single-direction ablation.
+        alpha: Scaling factor for the projected component. 0.0 = full ablation,
+            1.0 = no-op. Can be overridden per-call via :meth:`apply`.
+
+    Example::
+
+        hook = AblationHook(policy_net, layer_idx=5, direction=probe.coef_)
+        with hook.apply(alpha=0.0):
+            rewards = run_episodes(model, env, n=100)
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        layer_idx: int,
+        W: np.ndarray | None = None,
+        alpha: float = 0.0,
+        *,
+        direction: np.ndarray | None = None,
+    ) -> None:
+        # Accept either W or direction keyword
+        basis = W if W is not None else direction
+        if basis is None:
+            raise ValueError("Provide either W (subspace basis) or direction (1-D vector).")
+
+        device = next(model.parameters()).device
+        dtype = next(model.parameters()).dtype
+        basis_arr = np.asarray(basis, dtype=np.float32)
+        if basis_arr.ndim == 1:
+            # Single direction: normalise and reshape to (d, 1)
+            norm = float(np.linalg.norm(basis_arr))
+            if norm > 1e-9:
+                basis_arr = basis_arr / norm
+            basis_arr = basis_arr[:, None]
+        self._W = torch.tensor(basis_arr, dtype=dtype, device=device)  # (d_hidden, k)
+        self._alpha = alpha
+        self._model = model
+        self._layer_idx = layer_idx
+        self._hook = model[layer_idx].register_forward_hook(self._ablate)
+
+    def _ablate(
+        self,
+        module: nn.Module,
+        input: tuple,
+        output: torch.Tensor,
+    ) -> torch.Tensor:
+        """Forward hook: scale the subspace projection by alpha."""
+        proj = output @ self._W @ self._W.T
+        return output - proj + self._alpha * proj
+
+    def remove(self) -> None:
+        """Deregister the forward hook, restoring normal forward-pass behaviour."""
+        self._hook.remove()
+
+    @contextmanager
+    def apply(self, alpha: float) -> Generator[None, None, None]:
+        """Context manager that temporarily applies ablation at the given alpha.
+
+        Removes any existing hook, registers a new one with *alpha*, yields,
+        then restores the original hook.
+
+        Args:
+            alpha: Ablation strength. 0.0 = full ablation, 1.0 = no-op.
+
+        Yields:
+            None
+        """
+        self._hook.remove()
+        old_alpha = self._alpha
+        self._alpha = alpha
+        self._hook = self._model[self._layer_idx].register_forward_hook(self._ablate)
+        try:
+            yield
+        finally:
+            self._hook.remove()
+            self._alpha = old_alpha
+            self._hook = self._model[self._layer_idx].register_forward_hook(self._ablate)

rl_interrogate/env_wrappers.py

@@ -0,0 +1,147 @@
+"""
+rl_interrogate/env_wrappers.py
+-------------------------------
+Gymnasium environment wrappers for formation-flight policy evaluation.
+
+No Isaac Lab dependency — only imports gymnasium, numpy, and standard library.
+Importable from rl_interrogate without activating the Isaac Lab environment.
+
+Usage:
+    from rl_interrogate.env_wrappers import NoisyObservationWrapper
+
+    env = NoisyObservationWrapper(
+        base_env,
+        lateral_noise_std=0.05,
+        sensor_zero_prob=0.10,
+        rel_lead_delay=2,
+        seed=0,
+    )
+    obs, info = env.reset()
+    obs, reward, terminated, truncated, info = env.step(action)
+"""
+
+import numpy as np
+import gymnasium
+
+
+class NoisyObservationWrapper(gymnasium.Wrapper):
+    """
+    Adds sensor noise and delays to a formation-flight gymnasium environment.
+    NO Isaac Lab dependency — only imports gymnasium, numpy.
+
+    Three perturbations are applied to each observation:
+    1. Gaussian noise on the lateral position (index 15 per follower).
+    2. Random zeroing of the upwash sensor (index 12) with probability
+       ``sensor_zero_prob`` per step.
+    3. Temporal delay of rel_to_lead observations (indices 9–11) by
+       ``rel_lead_delay`` timesteps, using a circular buffer initialised to
+       zeros at reset.
+
+    Args:
+        env: The wrapped gymnasium environment.
+        lateral_noise_std: Std of Gaussian noise on lateral position obs (default 0.05).
+        sensor_zero_prob: Probability per step of zeroing upwash sensor index 12 (default 0.10).
+        rel_lead_delay: Number of timesteps to delay rel_to_lead obs indices 9-11 (default 2).
+        seed: RNG seed for reproducibility (default 0).
+    """
+
+    UPWASH_IDX: int = 12
+    REL_LEAD_SLICE: slice = slice(9, 12)
+    LATERAL_IDX: int = 15  # lateral y position index in per-follower obs
+
+    def __init__(
+        self,
+        env: gymnasium.Env,
+        lateral_noise_std: float = 0.05,
+        sensor_zero_prob: float = 0.10,
+        rel_lead_delay: int = 2,
+        seed: int = 0,
+    ) -> None:
+        super().__init__(env)
+        self.lateral_noise_std = lateral_noise_std
+        self.sensor_zero_prob = sensor_zero_prob
+        self.rel_lead_delay = rel_lead_delay
+        self._rng = np.random.default_rng(seed)
+
+        # Circular buffer — initialised properly in reset()
+        self._delay_buffer = np.zeros((self.rel_lead_delay, 3), dtype=np.float32)
+        self._buf_ptr: int = 0
+
+    # ------------------------------------------------------------------
+    # Public gymnasium.Wrapper interface
+    # ------------------------------------------------------------------
+
+    def reset(self, **kwargs):
+        """Reset env, zero delay buffer, apply noise to initial obs."""
+        # Zero the delay buffer and reset pointer
+        self._delay_buffer = np.zeros((self.rel_lead_delay, 3), dtype=np.float32)
+        self._buf_ptr = 0
+
+        result = self.env.reset(**kwargs)
+
+        # reset() may return (obs, info) or just obs depending on the env
+        if isinstance(result, tuple):
+            obs, *rest = result
+            obs = self._apply_noise(self._to_numpy(obs))
+            return (obs, *rest)
+        else:
+            obs = self._apply_noise(self._to_numpy(result))
+            return obs
+
+    def step(self, action):
+        """Step env, apply noise/delay to obs, return modified obs."""
+        result = self.env.step(action)
+
+        # step() returns (obs, reward, terminated, truncated, info)
+        obs, *rest = result
+        obs = self._apply_noise(self._to_numpy(obs))
+        return (obs, *rest)
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _to_numpy(self, obs) -> np.ndarray:
+        """Convert obs to a writable numpy float32 array if needed."""
+        try:
+            import torch
+            if isinstance(obs, torch.Tensor):
+                return obs.detach().cpu().numpy().astype(np.float32)
+        except ImportError:
+            pass
+        arr = np.asarray(obs, dtype=np.float32)
+        return arr.copy()  # ensure writeable
+
+    def _apply_noise(self, obs: np.ndarray) -> np.ndarray:
+        """
+        Apply all three perturbations to a single observation array.
+
+        Modifies ONLY indices 9–12 and index 15 (per follower).
+        All other indices are left unchanged.
+
+        Args:
+            obs: numpy float32 array of shape (obs_dim,) or (batch, obs_dim).
+
+        Returns:
+            Modified obs array (same shape and dtype).
+        """
+        # --- 1. Delay: read oldest entry from circular buffer ---
+        delayed = self._delay_buffer[self._buf_ptr].copy()
+
+        # --- 2. Write current rel_to_lead into the buffer slot ---
+        self._delay_buffer[self._buf_ptr] = obs[..., 9:12]
+
+        # --- 3. Advance pointer ---
+        self._buf_ptr = (self._buf_ptr + 1) % self.rel_lead_delay
+
+        # --- 4. Replace obs[9:12] with the delayed value ---
+        obs[..., 9:12] = delayed
+
+        # --- 5. Add Gaussian noise to lateral y (index 15) ---
+        obs[..., self.LATERAL_IDX] += self._rng.normal(0, self.lateral_noise_std)
+
+        # --- 6. Zero upwash sensor (index 12) with prob sensor_zero_prob ---
+        if self._rng.random() < self.sensor_zero_prob:
+            obs[..., self.UPWASH_IDX] = 0.0
+
+        return obs

rl_interrogate/llm_polarity/__init__.py

@@ -0,0 +1,35 @@
+"""LLM Polarity Inversion Validation Module
+
+This module implements cross-domain validation of the Polarity Inversion Discriminant (PID)
+applied to Large Language Models. The PID tests whether a feature is causally wired into
+a network's decision-making by inverting its polarity via hyperplane reflection and measuring
+the change in probe R².
+"""
+
+from rl_interrogate.llm_polarity.model_loader import ModelLoader
+from rl_interrogate.llm_polarity.dataset_builder import DatasetBuilder
+from rl_interrogate.llm_polarity.direction_extractor import DirectionExtractor
+from rl_interrogate.llm_polarity.probe_trainer import ProbeTrainer
+from rl_interrogate.llm_polarity.hyperplane_reflector import HyperplaneReflector
+from rl_interrogate.llm_polarity.output_analyzer import OutputAnalyzer
+from rl_interrogate.llm_polarity.control_runner import ControlRunner
+from rl_interrogate.llm_polarity.results_writer import ResultsWriter
+from rl_interrogate.llm_polarity.data_models import (
+    ExperimentResult,
+    ControlResult,
+    TransferResult,
+)
+
+__all__ = [
+    "ModelLoader",
+    "DatasetBuilder",
+    "DirectionExtractor",
+    "ProbeTrainer",
+    "HyperplaneReflector",
+    "OutputAnalyzer",
+    "ControlRunner",
+    "ResultsWriter",
+    "ExperimentResult",
+    "ControlResult",
+    "TransferResult",
+]

rl_interrogate/llm_polarity/control_runner.py

@@ -0,0 +1,203 @@
+"""ControlRunner component for executing control experiments to validate PID signal."""
+
+import numpy as np
+from typing import Dict
+from transformers import PreTrainedModel
+
+from .probe_trainer import ProbeTrainer
+from .hyperplane_reflector import HyperplaneReflector
+from .model_loader import ModelLoader
+
+
+class ControlRunner:
+    """Execute control experiments to validate PID signal.
+    
+    This class implements control experiments that test whether the PID signal
+    is spurious by using random directions, shuffled labels, and untrained models.
+    """
+    
+    @staticmethod
+    def run_random_direction_control(
+        model: PreTrainedModel,
+        layer_idx: int,
+        X: np.ndarray,
+        y: np.ndarray,
+        n_trials: int = 10,
+        alpha: float = 1.0,
+        test_size: float = 0.2,
+        random_state: int = 42,
+    ) -> Dict[str, float]:
+        """Run PID with random unit directions.
+        
+        Args:
+            model: Pre-trained language model (not used, kept for API consistency)
+            layer_idx: Layer index (not used, kept for API consistency)
+            X: Activations of shape (n_samples, d_model)
+            y: Labels of shape (n_samples,)
+            n_trials: Number of random direction trials
+            alpha: Ridge regression regularization parameter
+            test_size: Fraction of data for testing
+            random_state: Random seed for reproducibility
+            
+        Returns:
+            Dictionary with mean_delta_r2, std_delta_r2, mean_probe_r2
+        """
+        delta_r2_values = []
+        probe_r2_values = []
+        
+        np.random.seed(random_state)
+        
+        for trial in range(n_trials):
+            # Train probe
+            trainer = ProbeTrainer(alpha=alpha, test_size=test_size, random_state=random_state + trial)
+            trainer.fit(X, y)
+            
+            # Evaluate pre-inversion
+            train_r2, test_r2 = trainer.evaluate_test_set()
+            r2_pre = test_r2
+            probe_r2_values.append(r2_pre)
+            
+            # Generate random unit direction
+            d_model = X.shape[1]
+            random_direction = np.random.randn(d_model)
+            random_direction = random_direction / np.linalg.norm(random_direction)
+            
+            # Apply reflection to test set
+            X_test_inv = HyperplaneReflector.reflect(trainer.X_test, random_direction)
+            
+            # Evaluate post-inversion
+            r2_post = trainer.evaluate(X_test_inv, trainer.y_test)
+            
+            # Compute ΔR²
+            delta_r2 = r2_post - r2_pre
+            delta_r2_values.append(delta_r2)
+        
+        return {
+            'mean_delta_r2': float(np.mean(delta_r2_values)),
+            'std_delta_r2': float(np.std(delta_r2_values)),
+            'mean_probe_r2': float(np.mean(probe_r2_values)),
+        }
+    
+    @staticmethod
+    def run_shuffled_labels_control(
+        model: PreTrainedModel,
+        layer_idx: int,
+        X: np.ndarray,
+        y: np.ndarray,
+        n_trials: int = 10,
+        alpha: float = 1.0,
+        test_size: float = 0.2,
+        random_state: int = 42,
+    ) -> Dict[str, float]:
+        """Run PID with shuffled labels.
+        
+        Args:
+            model: Pre-trained language model (not used, kept for API consistency)
+            layer_idx: Layer index (not used, kept for API consistency)
+            X: Activations of shape (n_samples, d_model)
+            y: Labels of shape (n_samples,)
+            n_trials: Number of shuffled label trials
+            alpha: Ridge regression regularization parameter
+            test_size: Fraction of data for testing
+            random_state: Random seed for reproducibility
+            
+        Returns:
+            Dictionary with mean_delta_r2, std_delta_r2, mean_probe_r2
+        """
+        delta_r2_values = []
+        probe_r2_values = []
+        
+        np.random.seed(random_state)
+        
+        for trial in range(n_trials):
+            # Shuffle labels
+            y_shuffled = y.copy()
+            np.random.shuffle(y_shuffled)
+            
+            # Train probe with shuffled labels
+            trainer = ProbeTrainer(alpha=alpha, test_size=test_size, random_state=random_state + trial)
+            trainer.fit(X, y_shuffled)
+            
+            # Evaluate pre-inversion
+            train_r2, test_r2 = trainer.evaluate_test_set()
+            r2_pre = test_r2
+            probe_r2_values.append(r2_pre)
+            
+            # Extract probe direction
+            probe_direction = trainer.get_probe_direction()
+            
+            # Apply reflection to test set
+            X_test_inv = HyperplaneReflector.reflect(trainer.X_test, probe_direction)
+            
+            # Evaluate post-inversion
+            r2_post = trainer.evaluate(X_test_inv, trainer.y_test)
+            
+            # Compute ΔR²
+            delta_r2 = r2_post - r2_pre
+            delta_r2_values.append(delta_r2)
+        
+        return {
+            'mean_delta_r2': float(np.mean(delta_r2_values)),
+            'std_delta_r2': float(np.std(delta_r2_values)),
+            'mean_probe_r2': float(np.mean(probe_r2_values)),
+        }
+    
+    @staticmethod
+    def run_untrained_model_control(
+        model_name: str,
+        layer_idx: int,
+        X: np.ndarray,
+        y: np.ndarray,
+        alpha: float = 1.0,
+        test_size: float = 0.2,
+        random_state: int = 42,
+    ) -> Dict[str, float]:
+        """Run PID on randomly initialized model.
+        
+        Note: This control uses the provided activations X directly rather than
+        extracting from an untrained model, as the activations are already provided.
+        The key insight is that with random/untrained representations, the probe
+        should have low R² and ΔR² should be near zero.
+        
+        Args:
+            model_name: Model name (for logging purposes)
+            layer_idx: Layer index (for logging purposes)
+            X: Activations of shape (n_samples, d_model)
+            y: Labels of shape (n_samples,)
+            alpha: Ridge regression regularization parameter
+            test_size: Fraction of data for testing
+            random_state: Random seed for reproducibility
+            
+        Returns:
+            Dictionary with mean_delta_r2, std_delta_r2, mean_probe_r2
+        """
+        # For untrained model control, we use random activations
+        # to simulate what an untrained model would produce
+        np.random.seed(random_state)
+        X_random = np.random.randn(*X.shape)
+        
+        # Train probe on random activations
+        trainer = ProbeTrainer(alpha=alpha, test_size=test_size, random_state=random_state)
+        trainer.fit(X_random, y)
+        
+        # Evaluate pre-inversion
+        train_r2, test_r2 = trainer.evaluate_test_set()
+        r2_pre = test_r2
+        
+        # Extract probe direction
+        probe_direction = trainer.get_probe_direction()
+        
+        # Apply reflection to test set
+        X_test_inv = HyperplaneReflector.reflect(trainer.X_test, probe_direction)
+        
+        # Evaluate post-inversion
+        r2_post = trainer.evaluate(X_test_inv, trainer.y_test)
+        
+        # Compute ΔR²
+        delta_r2 = r2_post - r2_pre
+        
+        return {
+            'mean_delta_r2': float(delta_r2),
+            'std_delta_r2': 0.0,  # Single trial
+            'mean_probe_r2': float(r2_pre),
+        }

rl_interrogate/llm_polarity/data_models.py

@@ -0,0 +1,86 @@
+"""Data models for LLM Polarity Inversion Validation experiments."""
+
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class ExperimentResult:
+    """Results from a single PID experiment.
+    
+    Attributes:
+        model_name: Name of the model (e.g., "gpt2", "gpt2-medium", "llama-3.2-1b")
+        direction: Semantic direction (e.g., "sentiment", "factual", "odd_even", "toxicity")
+        layer_idx: Absolute layer index
+        layer_depth_pct: Layer depth as percentage (0.0, 0.25, 0.50, 0.75, 1.0)
+        probe_r2: R² before inversion
+        probe_accuracy: Classification accuracy before inversion
+        post_inversion_r2: R² after inversion
+        post_inversion_accuracy: Classification accuracy after inversion
+        delta_r2: R²_post - R²_pre
+        top_token_change_rate: Fraction of inputs with changed top token
+        kl_divergence: Mean KL divergence
+        perplexity_delta: Perplexity change
+        n_samples: Dataset size
+        runtime_seconds: Experiment duration
+    """
+    model_name: str
+    direction: str
+    layer_idx: int
+    layer_depth_pct: float
+    probe_r2: float
+    probe_accuracy: float
+    post_inversion_r2: float
+    post_inversion_accuracy: float
+    delta_r2: float
+    top_token_change_rate: float
+    kl_divergence: float
+    perplexity_delta: float
+    n_samples: int
+    runtime_seconds: float
+
+
+@dataclass
+class ControlResult:
+    """Results from a control experiment.
+    
+    Attributes:
+        control_type: Type of control ("random_direction", "shuffled_labels", "untrained")
+        model_name: Name of the model
+        direction: Semantic direction
+        layer_idx: Absolute layer index
+        mean_delta_r2: Mean ΔR² across trials
+        std_delta_r2: Standard deviation of ΔR² across trials
+        mean_probe_r2: Mean probe R² across trials
+        n_trials: Number of trials
+    """
+    control_type: str
+    model_name: str
+    direction: str
+    layer_idx: int
+    mean_delta_r2: float
+    std_delta_r2: float
+    mean_probe_r2: float
+    n_trials: int
+
+
+@dataclass
+class TransferResult:
+    """Results from cross-model transfer experiment.
+    
+    Attributes:
+        source_model: Model probe was trained on
+        target_model: Model probe was tested on
+        direction: Semantic direction
+        layer_idx: Absolute layer index (matched by depth percentage)
+        source_r2: R² on source model
+        transfer_r2: R² on target model
+        transfer_ratio: transfer_r2 / source_r2
+    """
+    source_model: str
+    target_model: str
+    direction: str
+    layer_idx: int
+    source_r2: float
+    transfer_r2: float
+    transfer_ratio: float