Author name cluster

Rahul Jain

Possible papers associated with this exact author name in Arrow. This page groups case-insensitive exact name matches and is not a full identity disambiguation profile.

14 papers

2 author rows

AAAI Conference 2025 Conference Paper

Markov Balance Satisfaction Improves Performance in Strictly Batch Offline Imitation Learning

Rishabh Agrawal
Nathan Dahlin
Rahul Jain
Ashutosh Nayyar

Imitation learning (IL) is notably effective for robotic tasks where directly programming behaviors or defining optimal control costs is challenging. In this work, we address a scenario where the imitator relies solely on observed behavior and cannot make environmental interactions during learning. It does not have additional supplementary datasets beyond the expert's dataset nor any information about the transition dynamics. Unlike state-of-the-art (SOTA) IL methods, this approach tackles the limitations of conventional IL by operating in a more constrained and realistic setting. Our method uses the Markov balance equation and introduces a novel conditional density estimation-based imitation learning framework. It employs conditional normalizing flows for transition dynamics estimation and aims at satisfying a balance equation for the environment. Through a series of numerical experiments on Classic Control and MuJoCo environments, we demonstrate consistently superior empirical performance compared to many SOTA IL algorithms.

PDF Details DOI

RLC Conference 2025 Conference Paper

Pure Exploration for Constrained Best Mixed Arm Identification with a Fixed Budget

Dengwang Tang
Rahul Jain
Ashutosh Nayyar
Pierluigi Nuzzo

We introduce the constrained best mixed arm identification (CBMAI) problem under unknown reward and costs wherein there are K arms, each of which is associated with a reward and multiple cost attributes. These are random, and come from distributions with unknown means. The best mixed arm is a probability distribution over a subset of the K arms that maximizes the expected reward while satisfying the expected cost constraints. We are specifically interested in a pure exploration problem under a fixed sampling budget with the goal of identifying the support of the best mixed arm. We propose a novel, parameter-free algorithm, called the Score Function-based Successive Reject (SFSR) algorithm, that combines the classical successive reject framework with a novel rejection criteria using a score function based on linear programming theory. We establish a performance guarantee for our algorithm by providing a theoretical upper bound on the probability of mis-identification of the support of the best mixed arm and show that it decays exponentially in the budget N and some constants that characterize the hardness of the problem instance. We also develop an information-theoretic lower bound on the error probability that shows that these constants appropriately characterize the problem difficulty. We validate this empirically on a number of problem instances.

PDF Details

RLJ Journal 2025 Journal Article

Pure Exploration for Constrained Best Mixed Arm Identification with a Fixed Budget

Dengwang Tang
Rahul Jain
Ashutosh Nayyar
Pierluigi Nuzzo

PDF Details

NeurIPS Conference 2025 Conference Paper

Robust LLM Alignment via Distributionally Robust Direct Preference Optimization

Zaiyan Xu
Sushil Vemuri
Kishan Panaganti
Dileep Kalathil
Rahul Jain
Deepak Ramachandran

A major challenge in aligning large language models (LLMs) with human preferences is the issue of distribution shift. LLM alignment algorithms rely on static preference datasets, assuming that they accurately represent real-world user preferences. However, user preferences vary significantly across geographical regions, demographics, linguistic patterns, and evolving cultural trends. This preference distribution shift leads to catastrophic alignment failures in many real-world applications. We address this problem using the principled framework of distributionally robust optimization, and develop two novel distributionally robust direct preference optimization (DPO) algorithms, namely, Wasserstein DPO (WDPO) and Kullback–Leibler DPO (KLDPO). We characterize the sample complexity of learning the optimal policy parameters for WDPO and KLDPO. Moreover, we propose scalable gradient descent-style learning algorithms by developing suitable approximations for the challenging minimax loss functions of WDPO and KLDPO. Our empirical experiments using benchmark data sets and LLMs demonstrate the superior performance of WDPO and KLDPO in substantially improving the alignment when there is a preference distribution shift.

PDF Details

NeurIPS Conference 2024 Conference Paper

e-COP : Episodic Constrained Optimization of Policies

Akhil Agnihotri
Rahul Jain
Deepak Ramachandran
Sahil Singla

In this paper, we present the e-COP algorithm, the first policy optimization algorithm for constrained Reinforcement Learning (RL) in episodic (finite horizon) settings. Such formulations are applicable when there are separate sets of optimization criteria and constraints on a system's behavior. We approach this problem by first establishing a policy difference lemma for the episodic setting, which provides the theoretical foundation for the algorithm. Then, we propose to combine a set of established and novel solution ideas to yield the e-COP algorithm that is easy to implement and numerically stable, and provide a theoretical guarantee on optimality under certain scaling assumptions. Through extensive empirical analysis using benchmarks in the Safety Gym suite, we show that our algorithm has similar or better performance than SoTA (non-episodic) algorithms adapted for the episodic setting. The scalability of the algorithm opens the door to its application in safety-constrained Reinforcement Learning from Human Feedback for Large Language or Diffusion Models.

PDF Details DOI

TCS Journal 2024 Journal Article

Space efficient algorithm for solving reachability using tree decomposition and separators

Rahul Jain
Raghunath Tewari

To solve reachability is to determine whether there is a path from one vertex to the other in a graph. Standard graph traversal algorithms such as DFS and BFS take linear time to solve reachability; however, their space complexity is also linear. On the other hand, Savitch's algorithm takes quasipolynomial time, although the space-bound is O ( log 2 ⁡ n ). In this paper, we study space-efficient algorithms for deciding reachability that runs in polynomial time. We show a polynomial-time algorithm that solves reachability in directed graphs using O ( w log ⁡ n ) space. Our algorithm requires access to a tree decomposition of width w for the underlying undirected graph of the input. This requirement can be waived for graphs for which recursive balanced vertex separators can be computed space-efficiently.

Details DOI

TMLR Journal 2023 Journal Article

Bridging Imitation and Online Reinforcement Learning: An Optimistic Tale

Botao Hao
Rahul Jain
Dengwang Tang
Zheng Wen

In this paper, we address the following problem: Given an offline demonstration dataset from an imperfect expert, what is the best way to leverage it to bootstrap online learning performance in MDPs. We first propose an Informed Posterior Sampling-based RL (iPSRL) algorithm that uses the offline dataset, and information about the expert's behavioral policy used to generate the offline dataset. Its cumulative Bayesian regret goes down to zero exponentially fast in $N$, the offline dataset size if the expert is competent enough. Since this algorithm is computationally impractical, we then propose the iRLSVI algorithm that can be seen as a combination of the RLSVI algorithm for online RL, and imitation learning. Our empirical results show that the proposed iRLSVI algorithm is able to achieve significant reduction in regret as compared to two baselines: no offline data, and offline dataset but used without suitably modeling the generative policy. Our algorithm can be seen as bridging online RL and imitation learning.

PDF Details

ICRA Conference 2023 Conference Paper

Learning Neuro-symbolic Programs for Language Guided Robot Manipulation

Namasivayam Kalithasan
Himanshu Singh 0002
Vishal Bindal
Arnav Tuli
Vishwajeet Agrawal
Rahul Jain
Parag Singla
Rohan Paul

Given a natural language instruction and an input scene, our goal is to train a model to output a manipulation program that can be executed by the robot. Prior approaches for this task possess one of the following limitations: (i) rely on hand-coded symbols for concepts limiting generalization beyond those seen during training [1] (ii) infer action sequences from instructions but require dense sub-goal supervision [2] or (iii) lack semantics required for deeper object-centric reasoning inherent in interpreting complex instructions [3]. In contrast, our approach can handle linguistic as well as perceptual variations, end-to-end trainable and requires no intermediate supervision. The proposed model uses symbolic reasoning constructs that operate on a latent neural object-centric representation, allowing for deeper reasoning over the input scene. Central to our approach is a modular structure consisting of a hierarchical instruction parser and an action simulator to learn disentangled action representations. Our experiments on a simulated environment with a 7-DOF manipulator, consisting of instructions with varying number of steps and scenes with different number of objects, demonstrate that our model is robust to such variations and significantly outperforms baselines, particularly in the generalization settings. The code, dataset and experiment videos are available at https://nsrmp.github.io

Details

TCS Journal 2023 Journal Article

Space-efficient algorithms for reachability in directed geometric graphs

Sujoy Bhore
Rahul Jain

The problem of graph Reachability is to decide whether there is a path from one vertex to another in a given graph. In this paper, we study the Reachability problem on three distinct graph families - intersection graphs of Jordan regions, unit contact disk graphs (penny graphs), and chordal graphs. For each of these graph families, we present space-efficient algorithms for the Reachability problem. For intersection graphs of Jordan regions, we show how to obtain a “good” vertex separator in a space-efficient manner and use it to solve the Reachability in polynomial time and O ( m 1 / 2 log ⁡ n ) space, where n is the number of Jordan regions, and m is the total number of crossings among the regions. We use a similar approach for chordal graphs and obtain a polynomial time and O ( m 1 / 2 log ⁡ n ) space algorithm, where n and m are the number of vertices and edges, respectively. However, for unit contact disk graphs (penny graphs), we use a more involved technique and obtain a better algorithm. We show that for every ϵ > 0, there exists a polynomial time algorithm that can solve Reachability in an n vertex directed penny graph, using O ( n 1 / 4 + ϵ ) space. We note that the method used to solve penny graphs does not extend naturally to the class of geometric intersection graphs that include arbitrary size cliques.

Details DOI

NeurIPS Conference 2022 Conference Paper

Matrix Multiplicative Weights Updates in Quantum Zero-Sum Games: Conservation Laws & Recurrence

Rahul Jain
Georgios Piliouras
Ryann Sim

Recent advances in quantum computing and in particular, the introduction of quantum GANs, have led to increased interest in quantum zero-sum game theory, extending the scope of learning algorithms for classical games into the quantum realm. In this paper, we focus on learning in quantum zero-sum games under Matrix Multiplicative Weights Update (a generalization of the multiplicative weights update method) and its continuous analogue, Quantum Replicator Dynamics. When each player selects their state according to quantum replicator dynamics, we show that the system exhibits conservation laws in a quantum-information theoretic sense. Moreover, we show that the system exhibits Poincare recurrence, meaning that almost all orbits return arbitrarily close to their initial conditions infinitely often. Our analysis generalizes previous results in the case of classical games.

PDF Details

AAAI Conference 2021 Conference Paper

A Sample-Efficient Algorithm for Episodic Finite-Horizon MDP with Constraints

Krishna C. Kalagarla
Rahul Jain
Pierluigi Nuzzo

Constrained Markov decision processes (CMDPs) formalize sequential decision-making problems whose objective is to minimize a cost function while satisfying constraints on various cost functions. In this paper, we consider the setting of episodic fixed-horizon CMDPs. We propose an online algorithm which leverages the linear programming formulation of repeated optimistic planning for finite-horizon CMDP to provide a probably approximately correctness (PAC) guarantee on the number of episodes needed to ensure a near optimal policy, i. e. , with resulting objective value close to that of the optimal value and satisfying the constraints within low tolerance, with high probability. The number of episodes needed is shown to have linear dependence on the sizes of the state and action spaces and quadratic dependence on the time horizon and an upper bound on the number of possible successor states for a state-action pair. Therefore, if the upper bound on the number of possible successor states is much smaller than the size of the state space, the number of needed episodes becomes linear in the sizes of the state and action spaces and quadratic in the time horizon.

PDF Details

NeurIPS Conference 2021 Conference Paper

Implicit Finite-Horizon Approximation and Efficient Optimal Algorithms for Stochastic Shortest Path

Liyu Chen
Mehdi Jafarnia-Jahromi
Rahul Jain
Haipeng Luo

We introduce a generic template for developing regret minimization algorithms in the Stochastic Shortest Path (SSP) model, which achieves minimax optimal regret as long as certain properties are ensured. The key of our analysis is a new technique called implicit finite-horizon approximation, which approximates the SSP model by a finite-horizon counterpart only in the analysis without explicit implementation. Using this template, we develop two new algorithms: the first one is model-free (the first in the literature to our knowledge) and minimax optimal under strictly positive costs; the second one is model-based and minimax optimal even with zero-cost state-action pairs, matching the best existing result from [Tarbouriech et al. , 2021b]. Importantly, both algorithms admit highly sparse updates, making them computationally more efficient than all existing algorithms. Moreover, both can be made completely parameter-free.

PDF Details

NeurIPS Conference 2017 Conference Paper

Learning Unknown Markov Decision Processes: A Thompson Sampling Approach

Yi Ouyang
Mukul Gagrani
Ashutosh Nayyar
Rahul Jain

We consider the problem of learning an unknown Markov Decision Process (MDP) that is weakly communicating in the infinite horizon setting. We propose a Thompson Sampling-based reinforcement learning algorithm with dynamic episodes (TSDE). At the beginning of each episode, the algorithm generates a sample from the posterior distribution over the unknown model parameters. It then follows the optimal stationary policy for the sampled model for the rest of the episode. The duration of each episode is dynamically determined by two stopping criteria. The first stopping criterion controls the growth rate of episode length. The second stopping criterion happens when the number of visits to any state-action pair is doubled. We establish $\tilde O(HS\sqrt{AT})$ bounds on expected regret under a Bayesian setting, where $S$ and $A$ are the sizes of the state and action spaces, $T$ is time, and $H$ is the bound of the span. This regret bound matches the best available bound for weakly communicating MDPs. Numerical results show it to perform better than existing algorithms for infinite horizon MDPs.

PDF Details

TCS Journal 2009 Journal Article

New bounds on classical and quantum one-way communication complexity

Rahul Jain
Shengyu Zhang

In this paper we provide new bounds on classical and quantum distributional communication complexity in the two-party, one-way model of communication. In the classical one-way model, our bound extends the well known upper bound of Kremer, Nisan and Ron [I. Kremer, N. Nisan, D. Ron, On randomized one-round communication complexity, in: Proceedings of The 27th ACM Symposium on Theory of Computing, STOC, 1995, pp. 596–605] to include non-product distributions. Let ϵ ∈ ( 0, 1 / 2 ) be a constant. We show that for a boolean function f: X × Y → { 0, 1 } and a non-product distribution μ on X × Y, D ϵ 1, μ ( f ) = O ( ( I ( X: Y ) + 1 ) ⋅ VC ( f ) ), where D ϵ 1, μ ( f ) represents the one-way distributional communication complexity of f with error at most ϵ under μ; VC ( f ) represents the Vapnik–Chervonenkis dimension of f and I ( X: Y ) represents the mutual information, under μ, between the random inputs of the two parties. For a non-boolean function f: X × Y → { 1, …, k } ( k ≥ 2 an integer), we show a similar upper bound on D ϵ 1, μ ( f ) in terms of k, I ( X: Y ) and the pseudo-dimension of f ′ = def f k, a generalization of the VC -dimension for non-boolean functions. In the quantum one-way model we provide a lower bound on the distributional communication complexity, under product distributions, of a function f, in terms of the well studied complexity measure of f referred to as the rectangle bound or the corruption bound of f. We show for a non-boolean total function f: X × Y → Z and a product distribution μ on X × Y, Q ϵ 3 / 8 1, μ ( f ) = Ω ( rec ϵ 1, μ ( f ) ), where Q ϵ 3 / 8 1, μ ( f ) represents the quantum one-way distributional communication complexity of f with error at most ϵ 3 / 8 under μ and rec ϵ 1, μ ( f ) represents the one-way rectangle bound of f with error at most ϵ under μ. Similarly for a non-boolean partial function f: X × Y → Z ∪ { ∗ } and a product distribution μ on X × Y, we show, Q ϵ 6 / ( 2 ⋅ 1 5 4 ) 1, μ ( f ) = Ω ( rec ϵ 1, μ ( f ) ).

Details DOI