Tabitha Edith Lee

Gated Transformer World Models for Model-Based Reinforcement Learning

Model-based reinforcement learning has emerged as a promising approach for training sample-efficient policies by learning compact latent dyn… (see more)amics models that simulate, plan, and act through imagination. Recent work has shown that transformer-based world models outperform recurrent architectures due to their ability to capture long-range temporal dependencies. However, standard self-attention mechanisms apply the attention score uniformly across all dimensions of the token's representation, lacking an explicit mechanism to regulate information flow across dimensions. We introduce a gating mechanism for transformer-based world models that implements a feature-wise modulation mechanism within the attention computation. The gating module adaptively scales dimensions in a data-dependent way, enabling the model to selectively control contextual information. Our gating mechanism achieves greater sample-efficient policy learning through normalized mean performance improvements for the Atari 100k benchmark when applied to the STORM and TWISTER transformer world models.

2026-10-02

WM_PAI @ Advances in Neural Information Processing Systems (poster)

openreview.net

CURATE: Automatic Curriculum Learning for Reinforcement Learning Agents through Competence-Based Curriculum Policy Search in Structured Task Spaces

Tabitha Edith Lee

Nan Rosemary Ke

Sarvesh Patil

Annya Dahmani

Eunice Yiu

Esra'a Saleh

Alison Gopnik

Oliver Kroemer

Glen Berseth

Due to fundamental exploration challenges without informed priors or specialized algorithms, agents may be unable to consistently receive in… (see more)formative rewards, leading to inefficient or intractable learning. To address these challenges, we introduce CURATE, an automatic curriculum learning algorithm for reinforcement learning agents in structured task spaces of monotonic difficulty. Through "exploration by exploitation," CURATE dynamically scales the task difficulty to match the agent's current competence. By exploiting its current capabilities that were learned in easier tasks, the agent improves its exploration in more difficult tasks. Our key insight is that the learning improvement in tasks that are close to those used for training is inversely proportional to their difficulty, and an agent that chooses a nearby distribution of the easiest unsolved tasks at any given time can automatically induce an easiest-to-hardest curriculum in these task spaces. To achieve this, CURATE conducts policy search in the task space to learn the best task distribution for training. As the agent's mastery grows, the learned curriculum adapts in an approximately easiest-to-hardest and task-directed fashion, efficiently culminating in a performant agent. Our experiments across three diverse domains (MiniGrid, Procgen, BipedalWalker) demonstrate that CURATE learns effective curricula for sample efficiency and proficiency with the potential for yielding broadly capable agents, matching or exceeding prior curriculum methods that do not require informed initialization or predefined schedules.

2026-08-30

Transactions on Machine Learning Research (accepted)

openreview.net

AI Agent Safety is a Reinforcement Learning Problem

Reginald McLean

Tabitha Edith Lee

Montaser Mohammedalamen

Kevin Roice

Glen Berseth

Patrick M. Pilarski

Marlos C. Machado

Alyssa Lefaivre Škopac

Benjamin Rosman

With the rapid advancement and deployment of Agentic AI, our scientific understanding of capabilities and limitations has not kept pace, lea… (see more)ding to cases where AI agents cause harm. We argue that many of these safety limitations are not novel problems. Instead, the safety challenges currently facing AI agents can be seen as instances of problems the reinforcement learning (RL) community has studied rigorously for decades. The core of this argument concerns the problem formulation of AI agents. AI agents are designed to solve sequential decision-making problems: problems with long-term objectives in which actions have delayed consequences. To model these types of problem, the problem is set up the problem such that the agent receives observations, feedback on its progress, and then takes actions. This is precisely the formulation of the RL problem. In this paper, we formalize the problem equivalence, which we then leverage to argue that \textbf{AI Agent safety is a reinforcement learning problem: the failure modes currently observed in deployed AI agents are structural instances of problems RL has formalized for decades, and the RL safety literature provides principled tools to diagnose and address them.}. We conclude with a call for deliberate collaboration between the RL and AI agent research communities: AI agent researchers gain access to principled frameworks, while RL researchers gain a class of real-world problems that could expose fundamental gaps in current RL benchmarks and theory.

2026-05-22

AIWILD @ International Conference on Machine Learning (published)

openreview.net