Technical Summary: ATLAS (Active Theory Learning for Automated Science)

Problem Statement

Advancing scientific understanding in cognitive science and related fields requires identifying the mechanistic algorithms that govern behavior. This process traditionally relies on an iterative, closed-loop cycle of hypothesis generation (computational modeling) and hypothesis-driven data collection (experimentation). However, this cycle is often slow and resource-intensive because experiments are costly.

Existing data-driven model discovery methods can identify interpretable models but typically require large datasets, making them impractical for scenarios where data collection is expensive. Conversely, Active Learning (AL) and Optimal Experiment Design (OED) methods exist to maximize information gain, but they often optimize for black-box predictive models (lacking interpretability) or discriminate among pre-specified mechanistic models (risking misspecification and limiting discovery).

The core challenge addressed by this paper is to create a framework that automates the discovery of interpretable mechanistic models while maximizing sample efficiency. The goal is to design experiments that optimally distinguish between competing hypotheses about an agent's internal computational structure using minimal data.

Methodology: The ATLAS Pipeline

The authors introduce ATLAS (Active Theory Learning for Automated Science), an active learning framework that iterates between generating diverse mechanistic hypotheses and designing experiments to distinguish them. The pipeline consists of three main components operating in cycles:

1. Hypothesis Generator (Ensemble of Disentangled RNNs)

Instead of a fixed set of models, ATLAS maintains an ensemble of competing hypotheses instantiated as Disentangled Recurrent Neural Networks (DisRNNs).

• Structure: DisRNNs utilize information bottlenecks to encourage solutions with a small number of sparsely interacting latent ("cognitive") variables.
• Interpretability: The interactions between these latent variables can be interpreted as a computational graph representing the underlying mechanism.
• Diversity: The generator trains a large set of networks with varying hyperparameters (specifically the sparsity penalty) and random seeds. It selects a subset for the ensemble using a probabilistic sampling strategy (offset softmax) over cross-validation likelihoods. This ensures the ensemble contains diverse structural hypotheses, ranging from underfitting to overfitting, rather than just parameter variations of a single model.

2. Experiment Optimizer

The optimizer seeks an experiment design $D$ that maximizes the Expected Information Gain (EIG) regarding the identity of the true generative model.

• Disagreement Metric: The EIG is approximated by the disagreement among the ensemble members. Specifically, it calculates the difference between a model's self-agreement (log-likelihood of its own simulated trajectories) and its cross-agreement (log-likelihood of trajectories generated by other models).
• Optimization Strategy: The experiment design is constrained to a binary $T \times A$ matrix (defining rewards for $A$ actions over $T$ timesteps). The optimizer uses an iterative hill-climbing strategy to flip bits in this matrix to maximize the ensemble disagreement score. To avoid local optima, this process is repeated with 128 random initializations.
• Output: The result is a temporally structured sequence of rewards designed to force the ensemble members to diverge in their behavioral predictions.

3. Experiment Runner

The optimized experiment is executed on the ground-truth agent (the "environment"). The resulting behavioral trajectory is added to the growing dataset, which is then used to retrain and update the hypothesis ensemble in the next cycle.

Key Contributions

1. Framework Integration: The paper demonstrates that combining disagreement-based active learning with automatic discovery of interpretable mechanistic models creates a practical framework for closed-loop scientific discovery (ATLAS).
2. DisRNN Efficacy: It shows that Disentangled RNNs serve as effective hypothesis generators within ATLAS, surfacing sufficiently diverse structural hypotheses to enable strong experimental design.
3. Binary Matrix Optimizer: The authors develop an optimizer for binary experiment designs that successfully drives disagreement between handcrafted models and performs well within the ATLAS loop.
4. Application to RL Discovery: The framework is applied to discovering reinforcement learning algorithms (specifically Q-learning and Leaky Actor-Critic) from overt behavior, demonstrating that it can recover ground-truth structures with an order of magnitude fewer experiments than random baselines.

Results

The authors evaluated ATLAS on the task of recovering the computational structure of two ground-truth agents: a Q-learning agent and a Leaky Actor-Critic agent. Performance was measured against two baselines:

1. Random Experimentation: Independent and identically distributed (i.i.d.) rewards or Gaussian random walks.
2. Expert-Designed Experiments: Fixed tasks derived from literature (e.g., specific Gaussian random walks used in prior cognitive science studies).

Key Findings

• Sample Efficiency: ATLAS achieved a 5–10× improvement in sample efficiency across all metrics compared to random experimentation.
  • Behavioral Similarity: Models trained on ATLAS data matched the ground-truth agent's behavior on held-out tests significantly faster (e.g., after 20 experiments for Q-learning vs. 1,000 for random).
  • Structural Recovery: ATLAS successfully recovered the correct computational graph (isomorphism with the ground truth) for all 8 independent runs after 100 experiments. In contrast, random baselines required 1,000 experiments to reach similar performance for Q-learning and failed to recover the structure for the Actor-Critic agent even after 1,000 experiments.
  • Dynamical Similarity: Using "approximate bisimulation" (measuring the ability of models to simulate each other's internal states), ATLAS-trained models showed significantly lower error rates than random baselines after 100 experiments.
• Comparison to Experts: ATLAS matched or exceeded the performance of expert-designed experiments. For the Leaky Actor-Critic agent, ATLAS significantly outperformed the expert baseline in structural recovery and dynamical similarity.
• Experiment Structure: The experiments designed by ATLAS exhibited complex, non-intuitive temporal structures (e.g., specific patterns of reward omission and block switching) that were tailored to the specific learning rates and dynamics of the agents, differing significantly from standard researcher-designed tasks.

Significance and Claims

The paper claims that ATLAS represents a significant step toward accelerating human-interpretable insights in cognitive science. By automating the loop between data-driven hypothesis generation and hypothesis-driven experiment selection, ATLAS addresses the "rate-limiting step" of data collection in scientific inquiry.

The authors emphasize that their approach specifically targets interpretable mechanistic models rather than black-box predictors. This distinction is crucial for cognitive science, where the goal is to understand the underlying algorithms of behavior, not just predict it. The results suggest that ATLAS can uncover the true computational structure of agents using far fewer experiments than traditional methods, potentially reducing the time and cost of scientific discovery in domains where mechanistic models are essential.

The authors remain modest regarding the current scope, noting that the method is currently tested in in silico analogs with known ground truths. They acknowledge that training ensembles from scratch at each step is a computational bottleneck and suggest future work could involve fine-tuning approaches, optimizing batches of experiments, or applying the framework to real-world data collection and domains beyond reinforcement learning.