Dense and Diverse Goal Coverage in Multi Goal Reinforcement Learning
Sagalpreet Singh, Rishi Saket, Aravindan Raghuveer
TL;DR
The paper addresses multi-goal reinforcement learning by introducing a policy-mixture framework that jointly maximizes return and disperses visitation across goal states, formalized through ${\mathcal{Z}}({\overline{\pi}}) = J_\gamma({\overline{\pi}}) + {\mathcal{I}}^{{\mathcal{S}}^+}(d[{\overline{\pi}}])$ and optimized via a Frank–Wolfe-based DDGC algorithm. At each iteration, a batch RL subroutine (FQI) learns a new policy from mixture-derived rewards, progressively expanding the mixture to cover more goals while preserving high return; theoretical guarantees bound the sub-optimality with respect to the objective. The approach is extended to continuous spaces using FQI/FAC and practical extensions like exploratory sampling and a goal buffer to robustly discover and retain goal states. Empirical results on synthetic MDPs and Brax/JaxGCRL benchmarks demonstrate that DDGC achieves near-optimal return with substantially improved diversity over goal states, validating the value of explicitly balancing return and dispersion in goal coverage.
Abstract
Reinforcement Learning algorithms are primarily focused on learning a policy that maximizes expected return. As a result, the learned policy can exploit one or few reward sources. However, in many natural situations, it is desirable to learn a policy that induces a dispersed marginal state distribution over rewarding states, while maximizing the expected return which is typically tied to reaching a goal state. This aspect remains relatively unexplored. Existing techniques based on entropy regularization and intrinsic rewards use stochasticity for encouraging exploration to find an optimal policy which may not necessarily lead to dispersed marginal state distribution over rewarding states. Other RL algorithms which match a target distribution assume the latter to be available apriori. This may be infeasible in large scale systems where enumeration of all states is not possible and a state is determined to be a goal state only upon reaching it. We formalize the problem of maximizing the expected return while uniformly visiting the goal states as Multi Goal RL in which an oracle classifier over the state space determines the goal states. We propose a novel algorithm that learns a high-return policy mixture with marginal state distribution dispersed over the set of goal states. Our algorithm is based on optimizing a custom RL reward which is computed - based on the current policy mixture - at each iteration for a set of sampled trajectories. The latter are used via an offline RL algorithm to update the policy mixture. We prove performance guarantees for our algorithm, showing efficient convergence bounds for optimizing a natural objective which captures the expected return as well as the dispersion of the marginal state distribution over the goal states. We design and perform experiments on synthetic MDPs and standard RL environments to evaluate the effectiveness of our algorithm.