Portrait of Khimya Khetarpal

Khimya Khetarpal

Affiliate Member
Research Scientist, Google DeepMind
Research Topics
Machine Learning Theory
Online Learning
Reinforcement Learning
Representation Learning

Biography

Khimya Khetarpal is a Research Scientist at Google Deepmind. She earned her PhD in Computer Science from the Reasoning and Learning Lab at McGill University and Mila, advised by Doina Precup. She is broadly interested in artificial intelligence and reinforcement learning. Her current research interests focus on how RL agents learn to efficiently represent the world's knowledge, plan with it, and adapt to changes over time. Khimya’s work has appeared in leading AI journals and conferences including NeurIPS, ICML, AAAI, AISTATS, ICLR, The Knowledge Engineering Review, ACM, JAIR and TMLR. Her work has also been featured in MIT Technology Review. She was recognized as a TMLR expert reviewer in 2023, one of the Rising Stars in EECS 2020, a finalist for Three Minute Thesis (3MT) competition in AAAI 2019, selected for the Doctoral Consortium at AAAI 2019, and awarded Best Paper Award (3rd Price) for an ICML 2018 workshop on lifelong learning. Throughout her career, she has sought to actively mentor through initiatives such as co-founding the Mila peer advising initiative, teaching and assisting AI4Good Lab, volunteering at Skype A Scientist, and mentoring at FIRST Robotics.

Her research aims to (1) understand intelligent behavior that bridges both action and perception grounded in theoretical foundations of reinforcement learning, and (2) build AI agents to efficiently represent the world's knowledge, plan with it, and adapt to changes over time through learning and interaction.

She currently approaches this with the following research directions:

- Selective Attention for Fast Adaptation and Robustness

- Learning Abstractions and Affordances

- Discovery and Continual Reinforcement Learning

Current Students

Daniel Lawson
PhD - Université de Montréal
Principal supervisor :
Glen Berseth
daniel.lawson@mila.quebec
Website
Github
Google Scholar

Publications

Imagining Goals for Policy Generalization with Geometric Horizon Models
Daniel Lawson
Dhruv Sreenivas
Adriana Hugessen
Khimya Khetarpal
Glen Berseth
Training from offline data has allowed for substantial progress in domains such as robotics, leading to general-purpose policies that can be… (see more) easily applied zero-shot or efficiently finetuned for downstream tasks. However, training policies on offline data can lead to poor generalization, both due to the choice of modeling objective and from learning from a static dataset. In this work, we focus on the challenging task of zero-shot goal generalization, where a policy is evaluated on unseen tasks that require reusing its existing knowledge (compositional generalization). An avenue for improving a policy's generalization is through generating new experience through world models; however, such generation has proven difficult for longer horizons. Thus, to alleviate this issue, we propose TD-Aug, sampling from a geometric horizon model, which allows for directly imagining novel outcomes that can be achieved through composing existing knowledge. We demonstrate that training on these future outcomes as goals for goal-conditioned BC and offline RL policies improves generalization in stitching-based OGBench tasks.
2026-07-20
RLBRew @ Reinforcement Learning Conference (published)
openreview.net
openreview.net
Difference-Aware Retrieval Policies for Imitation Learning
Quinn Pfeifer
Ethan Pronovost
Paarth Shah
Khimya Khetarpal
Siddhartha Srinivasa
Abhishek Gupta
Parametric imitation learning via behavior cloning can suffer from poor generalization to out-of-distribution states due to compounding erro… (see more)rs during deployment. We show that reusing the training data during inference via a semi-parametric retrieval-based imitation learning approach can alleviate this challenge. We present Difference-Aware Retrieval Policies for Imitation Learning (DARP), a semi-parametric retrieval-based imitation learning approach that addresses this limitation by reparameterizing the imitation learning problem in terms of local neighborhood structure rather than direct state-to-action mappings. Instead of learning a global policy, DARP trains a model to predict actions based on
2026-06-07
arXiv (preprint)
doi.org
arxiv.org
Preventing Learning Stagnation in PPO by Scaling to 1 Million Parallel Environments
Michael Beukman
Khimya Khetarpal
Zeyu Zheng
Will Dabney
Jakob Foerster
Michael Dennis
Clare Lyle
Plateaus, where an agent's performance stagnates at a suboptimal level, are a common problem in deep on-policy RL. Focusing on PPO due to it… (see more)s widespread adoption, we show that plateaus in certain regimes arise not because of known exploration, capacity, or optimization challenges, but because sample-based estimates of the loss eventually become poor proxies for the true objective over the course of training. As a recap, PPO switches between sampling rollouts from several parallel environments online using the current policy (which we call the outer loop) and performing repeated minibatch SGD steps against this offline dataset (the inner loop). In our work we consider only the outer loop, and conceptually model it as stochastic optimization. The step size is then controlled by the regularization strength towards the previous policy and the gradient noise by the number of samples collected between policy update steps. This model predicts that performance will plateau at a suboptimal level if the outer step size is too large relative to the noise. Recasting PPO in this light makes it clear that there are two ways to address this particular type of learning stagnation: either reduce the step size or increase the number of samples collected between updates. We first validate the predictions of our model and investigate how hyperparameter choices influence the step size and update noise, concluding that increasing the number of parallel environments is a simple and robust way to reduce both factors. Next, we propose a recipe for how to co-scale the other hyperparameters when increasing parallelization, and show that incorrectly doing so can lead to severe performance degradation. Finally, we vastly outperform prior baselines in a complex open-ended domain by scaling PPO to more than 1M parallel environments, thereby enabling monotonic performance improvement up to one trillion transitions.
2026-03-05
Open MIND (preprint)
doi.org
arxiv.org
Affordances Enable Partial World Modeling with LLMs
Khimya Khetarpal
Gheorghe Comanici
Jonathan Richens
Jeremy Shar
Fei Xia
Laurent Orseau
Aleksandra Faust
Doina Precup
2026-02-10
ArXiv (preprint)
arxiv.org
arxiv.org
Robust Intervention Learning from Emergency Stop Interventions
Ethan Pronovost
Khimya Khetarpal
Siddhartha Srinivasa
2026-02-02
ArXiv (preprint)
arxiv.org
arxiv.org
Plasticity as the Mirror of Empowerment
David Abel
Michael Bowling
Andre Barreto
Will Dabney
Shi Dong
Steven Hansen
Anna Harutyunyan
Khimya Khetarpal
Clare Lyle
Razvan Pascanu
Georgios Piliouras
Doina Precup
Jonathan Richens
Mark Rowland
Tom Schaul
Satinder Singh
Agents are minimally entities that are influenced by their past observations and act to influence future observations. This latter capacity … (see more)is captured by empowerment, which has served as a vital framing concept across artificial intelligence and cognitive science. This former capacity, however, is equally foundational: In what ways, and to what extent, can an agent be influenced by what it observes? In this paper, we ground this concept in a universal agent-centric measure that we refer to as plasticity, and reveal a fundamental connection to empowerment. Following a set of desiderata on a suitable definition, we define plasticity using a new information-theoretic quantity we call the generalized directed information. We show that this new quantity strictly generalizes the directed information introduced by Massey (1990) while preserving all of its desirable properties. Under this definition, we find that plasticity is well thought of as the mirror of empowerment: The two concepts are defined using the same measure, with only the direction of influence reversed. Our main result establishes a tension between the plasticity and empowerment of an agent, suggesting that agent design needs to be mindful of both characteristics. We explore the implications of these findings, and suggest that plasticity, empowerment, and their relationship are essential to understanding agency
2025-12-02
Conference on Neural Information Processing Systems (Accept (spotlight))
doi.org
openreview.net
Long Range Navigator (LRN): Extending robot planning horizons beyond metric maps
Matt Schmittle
Rohan Baijal
Nathan Hatch
Rosario Scalise
Mateo Guaman Castro
Sidharth Talia
Khimya Khetarpal
Byron Boots
Siddhartha Srinivasa
2025-10-06
Proceedings of The 8th Conference on Robot Learning (published)
doi.org
proceedings.mlr.press
Self-Predictive Representations for Combinatorial Generalization in Behavioral Cloning
Daniel Lawson
Adriana Hugessen
Charlotte Cloutier
Glen Berseth
Khimya Khetarpal
Behavioral cloning (BC) methods trained with supervised learning (SL) are an effective way to learn policies from human demonstrations in do… (see more)mains like robotics. Goal-conditioning these policies enables a single generalist policy to capture diverse behaviors contained within an offline dataset. While goal-conditioned behavior cloning (GCBC) methods can perform well on in-distribution training tasks, they do not necessarily generalize zero-shot to tasks that require conditioning on novel state-goal pairs, i.e. combinatorial generalization. In part, this limitation can be attributed to a lack of temporal consistency in the state representation learned by BC; if temporally related states are encoded to similar latent representations, then the out-of-distribution gap for novel state-goal pairs would be reduced. Hence, encouraging this temporal consistency in the representation space should facilitate combinatorial generalization. Successor representations, which encode the distribution of future states visited from the current state, nicely encapsulate this property. However, previous methods for learning successor representations have relied on contrastive samples, temporal-difference (TD) learning, or both. In this work, we propose a simple yet effective representation learning objective,
2025-06-30
rl-conference.cc/RLC/2025/Workshop/RLBrew (published)
doi.org
openreview.net
Representation Learning via Non-Contrastive Mutual Information
Zhaohan Daniel Guo
Bernardo Avila Pires
Khimya Khetarpal
Dale Schuurmans
Bo Dai
2025-04-22
ArXiv (preprint)
doi.org
arxiv.org
A Unifying Framework for Action-Conditional Self-Predictive Reinforcement Learning
Khimya Khetarpal
Zhaohan Daniel Guo
Bernardo Avila Pires
Yunhao Tang
Clare Lyle
Mark Rowland
Nicolas Heess
Diana Borsa
Arthur Guez
Will Dabney
2025-04-22
Proceedings of The 28th International Conference on Artificial Intelligence and Statistics (published)
doi.org
proceedings.mlr.press
Cracking the Code of Action: A Generative Approach to Affordances for Reinforcement Learning
Lynn Cherif
Flemming Kondrup
David Venuto
Ankit Anand
Doina Precup
Khimya Khetarpal
Agents that can autonomously navigate the web through a graphical user interface (GUI) using a unified action space (e.g., mouse and keyboar… (see more)d actions) can require very large amounts of domain-specific expert demonstrations to achieve good performance. Low sample efficiency is often exacerbated in sparse-reward and large-action-space environments, such as a web GUI, where only a few actions are relevant in any given situation. In this work, we consider the low-data regime, with limited or no access to expert behavior. To enable sample-efficient learning, we explore the effect of constraining the action space through
2025-03-04
ICLR.cc/2025/Workshop/DL4C (published)
doi.org
openreview.net
Disentangling the Causes of Plasticity Loss in Neural Networks
Clare Lyle
Zeyu Zheng
Khimya Khetarpal
Hado van Hasselt
Razvan Pascanu
James Martens
Will Dabney
2025-02-16
Proceedings of The 3rd Conference on Lifelong Learning Agents (published)
doi.org
proceedings.mlr.press