Shagun Sodhani — accepted AI-conference papers

2025

AI Research Agents for Machine Learning: Search, Exploration, and Generalization in MLE-bench

Edan Toledo, Karen Hambardzumyan, Martin Josifoski, RISHI HAZRA, Nicolas Baldwin, Alexis Audran-Reiss +19

NeurIPS 2025spotlight

AI research agents are demonstrating great potential to accelerate scientific progress by automating the design, implementation, and training of machine learning models. We focus on methods for improving agents' performance on MLE-bench, a challenging benchmark where agents compete in Kaggle competi…

Cited by 0 Source ↗code ↗Scholar ↗

2025

MaestroMotif: Skill Design from Artificial Intelligence Feedback

Martin Klissarov, Mikael Henaff, Roberta Raileanu, Shagun Sodhani, Pascal Vincent, Amy Zhang +4

ICLR 2025oral

Describing skills in natural language has the potential to provide an accessible way to inject human knowledge about decision-making into an AI system. We present MaestroMotif, a method for AI-assisted skill design, which yields high-performing and adaptable agents. MaestroMotif leverages the capabi…

Cited by 2 Source ↗PDF ↗Scholar ↗

2025

The Automated LLM Speedrunning Benchmark: Reproducing NanoGPT Improvements

Bingchen Zhao, Despoina Magka, Minqi Jiang, Xian Li, Roberta Raileanu, Tatiana Shavrina +17

NeurIPS 2025poster

Rapidly improving large language models (LLMs) have the potential to assist in scientific progress. One critical skill in this endeavor is the ability to faithfully reproduce existing work. To evaluate the capability of AI agents to reproduce complex code in an active research area, we introduce the…

Cited by 0 Source ↗code ↗Scholar ↗

2024

Do Large Language Models Know How Much They Know?

Gabriele Prato, Jerry Huang, Prasanna Parthasarathi, Shagun Sodhani, Sarath Chandar

EMNLP 2024main

Large Language Models (LLMs) have emerged as highly capable systems and are increasingly being integrated into various uses. Nevertheless, the rapid advancement in their deployment trails a comprehensive understanding of their internal mechanisms, as well as a delineation of their capabilities and l…

Cited by 1 Source ↗PDF ↗Scholar ↗

2024

Harnessing small projectors and multiple views for efficient vision pretraining

Arna Ghosh, Kumar Krishna Agrawal, Shagun Sodhani, Adam Oberman, Blake Aaron Richards

NeurIPS 2024poster

Recent progress in self-supervised (SSL) visual representation learning has led to the development of several different proposed frameworks that rely on augmentations of images but use different loss functions. However, there are few theoretically grounded principles to guide practice, so practical…

Cited by 0 Source ↗PDF ↗code ↗Scholar ↗

2024

Motif: Intrinsic Motivation from Artificial Intelligence Feedback

Martin Klissarov, Pierluca D'Oro, Shagun Sodhani, Roberta Raileanu, Pierre-Luc Bacon, Pascal Vincent +2

ICLR 2024poster

Exploring rich environments and evaluating one's actions without prior knowledge is immensely challenging. In this paper, we propose Motif, a general method to interface such prior knowledge from a Large Language Model (LLM) with an agent. Motif is based on the idea of grounding LLMs for decision-ma…

Cited by 64 Source ↗PDF ↗code ↗Scholar ↗

2024

TorchRL: A data-driven decision-making library for PyTorch

Albert Bou, Matteo Bettini, Sebastian Dittert, Vikash Kumar, Shagun Sodhani, Xiaomeng Yang +2

ICLR 2024spotlight

PyTorch has ascended as a premier machine learning framework, yet it lacks a native and comprehensive library for decision and control tasks suitable for large development teams dealing with complex real-world data and environments. To address this issue, we propose TorchRL, a generalistic control l…

Cited by 46 Source ↗PDF ↗Scholar ↗

2024

When should we prefer Decision Transformers for Offline Reinforcement Learning?

Prajjwal Bhargava, Rohan Chitnis, Alborz Geramifard, Shagun Sodhani, Amy Zhang

ICLR 2024poster

Offline reinforcement learning (RL) allows agents to learn effective, return-maximizing policies from a static dataset. Three popular algorithms for offline RL are Conservative Q-Learning (CQL), Behavior Cloning (BC), and Decision Transformer (DT), from the class of Q-Learning, Imitation Learning, a…

Cited by 10 Source ↗PDF ↗code ↗Scholar ↗

2023

EpiK-Eval: Evaluation for Language Models as Epistemic Models

Gabriele Prato, Jerry Huang, Prasanna Parthasarathi, Shagun Sodhani, Sarath Chandar

EMNLP 2023long main

In the age of artificial intelligence, the role of large language models (LLMs) is becoming increasingly central. Despite their growing prevalence, their capacity to consolidate knowledge from different training documents—a crucial ability in numerous applications—remains unexplored. This paper pres…

Cited by 0 Source ↗code ↗Scholar ↗

2023

VIP: Towards Universal Visual Reward and Representation via Value-Implicit Pre-Training

Yecheng Jason Ma, Shagun Sodhani, Dinesh Jayaraman, Osbert Bastani, Vikash Kumar, Amy Zhang

ICLR 2023top-25%

Reward and representation learning are two long-standing challenges for learning an expanding set of robot manipulation skills from sensory observations. Given the inherent cost and scarcity of in-domain, task-specific robot data, learning from large, diverse, offline human videos has emerged as a p…

Cited by 295 Source ↗PDF ↗code ↗Scholar ↗

2022

Robust Policy Learning over Multiple Uncertainty Sets

Annie Xie, Shagun Sodhani, Chelsea Finn, Joelle Pineau, Amy Zhang

ICML 2022spotlight

Reinforcement learning (RL) agents need to be robust to variations in safety-critical environments. While system identification methods provide a way to infer the variation from online experience, they can fail in settings where fast identification is not possible. Another dominant approach is robus…

Cited by 18 Source ↗PDF ↗Scholar ↗

2022

The Neural Race Reduction: Dynamics of Abstraction in Gated Networks

Andrew Saxe, Shagun Sodhani, Sam Jay Lewallen

ICML 2022spotlight

Our theoretical understanding of deep learning has not kept pace with its empirical success. While network architecture is known to be critical, we do not yet understand its effect on learned representations and network behavior, or how this architecture should reflect task structure.In this work, w…

Cited by 38 Source ↗PDF ↗code ↗Scholar ↗

2021

IIRC: Incremental Implicitly-Refined Classification

Mohamed Abdelsalam, Mojtaba Faramarzi, Shagun Sodhani, Sarath Chandar

CVPR 2021poster

We introduce the 'Incremental Implicitly-Refined Classification (IIRC)' setup, an extension to the class incremental learning setup where the incoming batches of classes have two granularity levels. i.e., each sample could have a high-level (coarse) label like 'bear' and a low-level (fine) label lik…

Cited by 39 PDF ↗Scholar ↗

2021

Learning Robust State Abstractions for Hidden-Parameter Block MDPs

Amy Zhang, Shagun Sodhani, Khimya Khetarpal, Joelle Pineau

ICLR 2021poster

Many control tasks exhibit similar dynamics that can be modeled as having common latent structure. Hidden-Parameter Markov Decision Processes (HiP-MDPs) explicitly model this structure to improve sample efficiency in multi-task settings. However, this setting makes strong assumptions on the observab…

Cited by 57 Source ↗PDF ↗Scholar ↗

2021

Multi-Task Reinforcement Learning with Context-based Representations

Shagun Sodhani, Amy Zhang, Joelle Pineau

ICML 2021spotlight

https://drive.google.com/file/d/1lRV72XaKoxZjgQrLXBJhsM82x54_1Vc4/view?usp=sharing

Cited by 242 Source ↗PDF ↗code ↗Scholar ↗

2021

Recurrent Independent Mechanisms

Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio +1

ICLR 2021spotlight

We explore the hypothesis that learning modular structures which reflect the dynamics of the environment can lead to better generalization and robustness to changes that only affect a few of the underlying causes. We propose Recurrent Independent Mechanisms (RIMs), a new recurrent architecture in wh…

Cited by 389 Source ↗PDF ↗Scholar ↗

2020

Invariant Causal Prediction for Block MDPs

Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau +2

ICML 2020poster

Generalization across environments is critical to the successful application of reinforcement learning (RL) algorithms to real-world challenges. In this work we propose a method for learning state abstractions which generalize to novel observation distributions in the multi-environment RL setting. W…

Cited by 169 Source ↗PDF ↗code ↗Scholar ↗

2020

Reinforcement Learning with Competitive Ensembles of Information-Constrained Primitives

Anirudh Goyal, Shagun Sodhani, Jonathan Binas, Xue Bin Peng, Sergey Levine, Yoshua Bengio

ICLR 2020poster

Reinforcement learning agents that operate in diverse and complex environments can benefit from the structured decomposition of their behavior. Often, this is addressed in the context of hierarchical reinforcement learning, where the aim is to decompose a policy into lower-level primitives or option…

Cited by 55 Source ↗Scholar ↗