Source-linked AI summary
Reward-Free Exploration for Reinforcement Learning
Chi Jin, Akshay Krishnamurthy, Max Simchowitz, Tiancheng Yu
TL;DR
Reinforcement-learning exploration can be prohibitively expensive when rewards are unknown or repeatedly redesigned. The paper introduces reward-free exploration, collecting a reusable dataset before rewards are specified, and obtains near-optimal sample complexity with planning through black-box solvers. Its scope remains tabular, with extensions to unobserved pre-specified rewards and function approximation left as future directions.
Problem
Naive exploration can have exponential sample complexity, while prior work gives limited guidance for obtaining coverage sufficient to optimize many reward functions after data collection.
Method
The paper separates reward-free exploration from planning, using significant-state coverage proportional to maximum visitation probability and allowing arbitrary Batch-RL planners.
Results
˜O(S^2Apoly(H)/ε^2) exploration episodes yield ε-suboptimal policies for arbitrarily many reward functions, with a nearly matching Ω(S^2AH^2/ε^2) lower bound.
Takeaways & Limitations
A single reward-independent dataset can support adaptive post hoc reward optimization without further data collection.
Takeaways & Limitations
The lower bound does not directly apply when a pre-specified reward is merely unobserved during exploration, and function-approximation extensions remain open.
Abstract
from arXiv · showhide
Exploration is widely regarded as one of the most challenging aspects of reinforcement learning (RL), with many naive approaches succumbing to exponential sample complexity. To isolate the challenges of exploration, we propose a new "reward-free RL" framework. In the exploration phase, the agent first collects trajectories from an MDP $\mathcal{M}$ without a pre-specified reward function. After exploration, it is tasked with computing near-optimal policies under for $\mathcal{M}$ for a collection of given reward functions. This framework is particularly suitable when there are many reward functions of interest, or when the reward function is shaped by an external agent to elicit desired behavior. We give an efficient algorithm that conducts $\tilde{\mathcal{O}}(S^2A\mathrm{poly}(H)/ε^2)$ episodes of exploration and returns $ε$-suboptimal policies for an arbitrary number of reward functions. We achieve this by finding exploratory policies that visit each "significant" state with probability proportional to its maximum visitation probability under any possible policy. Moreover, our planning procedure can be instantiated by any black-box approximate planner, such as value iteration or natural policy gradient. We also give a nearly-matching $Ω(S^2AH^2/ε^2)$ lower bound, demonstrating the near-optimality of our algorithm in this setting.
1 Introduction
The paper frames exploration without reward information as a way to collect reusable data for optimizing many possible rewards, and develops near-optimal coverage guarantees for this setting.
- Motivation: Randomized exploration can require exponential time to reach high-reward states, motivating strategies that deliberately seek informative states.The paper contrasts naive random walks with deliberately exploratory methods.
- Reward-free RL: Reward-free exploration asks how to explore efficiently without rewards, then compute near-optimal policies for any reward function after exploration.The exploration and planning phases are separated.
- Main contribution: The proposed algorithm uses a reward-independent dataset with sufficient coverage for standard Batch-RL solvers and supports all reward functions simultaneously.No additional data collection is needed when new reward functions are considered.
- Related work: The planning phase accommodates arbitrary Batch-RL solvers, including value iteration and natural policy gradient, while prior coverage approaches have suboptimal guarantees or incomplete end-to-end analyses.The comparison includes RMAX, function-approximation-derived methods, and Hazan et al.’s Frank-Wolfe approach.
- Sample complexity: The algorithm requires ˜O(S^2Apoly(H)/ε^2) exploration episodes, while the lower bound is Ω(S^2AH^2/ε^2), establishing near-optimality.The paper identifies an S-factor price for good coverage compared with pre-specified-reward RL.
- Technical novelty: The exploration method separates difficult-to-reach states with negligible reward-optimization contribution from significant states, which it visits with probability proportional to λ(s).Here λ(s) is the maximum visitation probability of state s under any policy.
2 Preliminaries
The paper studies tabular episodic MDPs and separates reward-free trajectory collection from later planning with a supplied reward function.
- MDP setting: The environment is a tabular episodic MDP with finite state and action sets, horizon H, time-dependent transitions, and rewards normalized to [0,1].The transition kernel specifies next-state distributions for each state, action, and step.
- Interaction protocol: In each episode, the agent observes a state, selects an action according to a possibly stochastic non-stationary policy, receives reward, and transitions until horizon H.The initial state is drawn from an unknown initial distribution.
- Optimality definitions: The value and action-value functions represent expected remaining rewards, and an optimal policy maximizes these quantities under the Bellman recursion.The objective is to find an ε-optimal policy.
- Exploration phase: During reward-free exploration, the agent follows trajectories without reward information and stores visited states, actions, and transitions in a dataset D.The dataset is collected over K episodes.
- Planning phase: After exploration, interaction stops and the agent uses D to compute a near-optimal policy for a potentially adversarially designed reward function.Performance depends on how many exploration episodes are needed for reliable planning.
3 Main Results
The paper separates reward-free exploration from later planning, collecting coverage that supports near-optimal policies for arbitrarily many reward functions. Its key mechanism focuses exploration on significant states and uses approximate planning over an empirical transition model.
- Main Results: O(H5S2A/ε2) episodes suffice for computing ε-optimal policies for arbitrarily many reward functions after exploration.The guarantee holds with high probability, up to logarithmic factors and lower-order terms.
- Main Results: The algorithm learns exploration policies, collects trajectories, estimates transitions, and invokes a planner separately for each reward function.Planning uses the empirical transition matrix and can accommodate arbitrary approximate MDP solvers.
- 3.1 Exploration Phase: A δ-significant state is reachable with probability greater than δ under some policy, while insignificant states can be neglected for near-optimal planning.This distinction addresses states that are difficult or impossible to reach under any policy.
- 3.1 Exploration Phase: The exploration distribution covers every significant state-action pair with probability proportional to its maximum visitation probability under any policy.The strategy first learns policies maximizing reachability for individual states, then combines them to form the exploration policy set.
- 3.2 Planning Phase: The planning guarantee transfers approximate optimality from the empirical MDP to the true MDP, yielding a 3ε-suboptimal policy simultaneously for any reward function.The transition estimate is accurate on significant states, while the total contribution of insignificant states is small.
- 3.3 Approximate MDP Solvers: Natural Policy Gradient is one supported approximate solver, updating policies using value estimates, exponential scaling, and normalization.Value iteration is also presented as a special case of the paper’s black-box planning interface.
4 Lower Bound
The lower-bound construction shows that reward-free exploration requires nearly quadratic dependence on the state count, in addition to action, horizon, and accuracy factors. It does so by encoding difficult transition-learning problems at one state and embedding many copies across a binary tree.
- Ω(H^2S^2A/ε^2) trajectories are necessary for the reward-free guarantee, even with randomized or history-dependent policies.The result holds under stated conditions on S, A, H, and ε, and even when rewards and transitions are identical across stages.
- The lower bound matches the upper bound in S, A, and ε up to logarithmic factors and lower-order terms.The comparison remains open when correctness is required only for a smaller, fixed set of reward vectors.
- 4.1 Reward Free Exploration at a Single State: At a single initial state, different actions encode distinct transition distributions over 2n absorbing states, forcing learning of Ω(n) information.The construction uses action-independent rewards, so the challenge comes from identifying transitions rather than reward variation.
- 4.1 Reward Free Exploration at a Single State: Learning near-optimal policies for exponentially many reward vectors yields an Ω(SA/ε^2) trajectory lower bound in the single-state construction.Appropriate reward vectors decode a packing of transition vectors, and a variant of Fano’s inequality supplies the information-theoretic argument.
- 4.2 Lower Bound for Multiple States: Embedding Ω(S) single-state instances in a binary tree forces the learner to estimate transition vectors at many reachable states.Varying the rewarded tree state requires learning n instances, giving Ω(S^2A/ε^2), while an absorbing-state reward chain adds the H^2 factor.
5 Conclusion
The conclusion frames reward-free RL as a two-phase process: collect trajectories without reward information, then plan for many reward functions without further interaction. The approach achieves near-optimal exploration complexity and supports black-box batch planners, while function approximation and partially unobserved rewards remain open directions.
- Reward-free RL separates reward-independent trajectory collection from post-exploration policy computation for a collection of reward functions.After exploration, the learner cannot interact with the MDP and instead plans under its transitions.
- The algorithm uses Õ(S^2Apoly(H)/ε^2) exploration episodes and returns ε-suboptimal policies for arbitrarily many adaptively chosen reward functions.The guarantee uses the exploration dataset alone, without additional data collection.
- Planning is compatible with arbitrary black-box approximate planners, including value iteration and natural policy gradient.ZERORMAX provides a concrete value-iteration instantiation after constructing an empirical MDP for each planning configuration.
- Future work includes reward-free exploration when a prespecified reward is unobserved during exploration and extensions to function approximation.The lower-bound proof does not directly cover the potentially easier prespecified-reward setting, and generalizing significant states and coverage remains open.
- ZERORMAX marks sufficiently sampled state-step-action pairs as known and explores through empirical MDPs updated across episodes.The known set contains states visited enough times to estimate the corresponding transition kernel; each episode computes a policy, samples a trajectory, and updates the set.
A.1 Analysis
The analysis formalizes escape from known states and compares three MDPs to transfer guarantees from the intermediate model to the true MDP. It then bounds the number of problematic episodes and derives the resulting sample complexity.
- Escape probability: Escape probability measures the chance of visiting unknown states and is central to analyzing ZERORMAX.The analysis argues that large escape probability cannot persist across many episodes because unknown states can be visited only a bounded number of times.
- Episode control: At most O(...) episodes can have large escape probability, with high probability.This is established through a lemma controlling episodes in which the policy visits unknown states frequently.
- MDP comparison: Three MDPs are compared: the true MDP M, the planning model ˆM_K, and the intermediate model M_K with self-loops on unknown states.The proof shows value closeness between ˆM_K and M_K, then between M_K and M, yielding the desired planning guarantee.
- Value transfer: For any policy and reward function, the values on the compared MDPs are close with high probability under the good-episode conditions.The argument separately bounds differences between M_K and ˆM_K and between M_K and M, then combines them.
- Sample complexity: The analyzed alternative sample complexity scales with ǫ^-3 and polynomially rather than logarithmically with 1/p.This scaling is identified as quite poor in the analysis.
B MaxEnt Exploration
This section analyzes maximum-entropy exploration as an alternative route to obtaining state coverage. It shows that the required optimization accuracy can be set to O(1/S), but the resulting sufficient sample bound is O(S5).
- MaxEnt objective: Maximum-entropy exploration seeks a policy whose state occupancy measure provides coverage comparable to the desired exploratory policy.The analysis contrasts an exact optimizer's coverage with the approximate optimizer produced by a Frank-Wolfe-style method.
- Approximate optimization: An optimization error of O(1/S) suffices for the maximum-entropy approach to obtain the needed coverage guarantee.The proof uses a mixture-policy argument to show how objective improvement depends on the error scale.
- Sample complexity: The resulting sufficient sample complexity scales with S5 after setting ε = 1/S.The analysis calls this scaling quite poor and does not claim it is necessary for the method.
- EULER-based exploration: The paper's exploration proof uses EULER and adapts its regret analysis to episodic state augmentation.The construction designs rewards that are zero except at one target state and applies the resulting guarantee across significant state-time pairs.
- Transition estimation: Empirical transition estimates are controlled uniformly over stages, deterministic policies, and bounded value functions using concentration and covering arguments.The proof invokes Bernstein's inequality, covers deterministic policies and value-related quantities, and applies a union bound over stages.
C.3 Proof of Theorem 3.1
The proof of Theorem 3.1 chooses the significant-state threshold and per-state exploration budget, then combines the coverage result with the main theorem. It also distinguishes the episodic convergence analysis for natural policy gradient.
- Proof of Theorem 3.1: δ = ǫ/(2SH2) and N0 ≥ cS3AH6ι3/ǫ are sufficient parameter choices in the proof.Because N0 episodes are required for each state-time pair, the construction contributes O(cS4AH7ι3/ǫ) episodes.
- Proof of Theorem 3.1: The total episode bound combines the cost of finding the exploratory policy with the first term supplied by Theorem 3.5.The proof concludes by adding these two contributions.
- Natural policy gradient: The episodic setting requires a different natural-policy-gradient guarantee because not every state is reachable at the first step.The section notes that the corresponding analysis from the tabular infinite-horizon setting does not directly apply.
- Natural policy gradient: The normalization constant in the natural policy-gradient update is not unique, and the proof chooses a form that simplifies the analysis.Different choices are stated to produce essentially the same algorithm.
- Natural policy gradient: Natural policy gradient uses a monotone value property to convert a cumulative-regret bound into a convergence rate for the final iteration.The proof applies the performance difference lemma and bounds the normalization term appearing in the regret analysis.
D Proof of Lower Bound
The lower-bound proof develops a formal environment, transition classes, and reward classes, then defines reward-free algorithms and their correctness. It permits arbitrary randomized history-dependent exploration policies.
- Proof strategy: The lower-bound construction proceeds from a single-state transition-learning lemma to an embedding across n states.The proof is organized through notation, the single-state lower bound, the multi-state embedding, and the final theorem.
- Formal setting: An environment specifies a finite state space, number of actions, and horizon, while transition and reward classes describe admissible dynamics and rewards.The transition class includes transition and initial-state distributions; the reward class contains functions from state-action pairs to [0,1].
- Algorithm definition: A reward-free MDP algorithm first collects a random number of trajectories without a reward and later returns policies for a sequence of reward vectors.The exploration and planning phases are separated in the formal definition.
- Correctness: Correctness requires ε-optimality for every reward vector in the supplied finite sequence with probability at least 1 − p.The guarantee is evaluated under the relevant transition operator and applies simultaneously across the requested rewards.
- Policy class: The lower bound allows exploration policies to be arbitrary randomized mappings of observed histories and a random seed.This broad policy class strengthens the scope of the lower-bound formulation.
D.2 Learning A Single Instance
The single-instance lower-bound construction uses one initial state, action-dependent near-uniform transitions to absorbing states, and action-independent rewards. Learning this class requires recovering substantial information about the transition probabilities.
- Ω(nA/ε^2) trajectories are required to learn transition probabilities at a given state.
- The hard instances start at state 0 and transition to one of 2n terminal states under an unknown, near-uniform distribution.The terminal states are absorbing, and the transition probabilities are ε/2n-close to uniform in ℓ∞.
- The environment has horizon 2 with state space {0, 1, ..., 2n}.
- Rewards assign zero to the initial state and action-independent values to terminal states.The reward vector can assign terminal-state values through ν.
- Lemma D.2 establishes a transition distribution under which any sufficiently accurate learner must satisfy the lower-bound trajectory requirement.Its conditions include ε ≤ 1, p ≤ 1/2, A ≥ 2, and n ≥ c0 log^2 A.
D.3 Learning Transitions at n states: Proof of Theorem 4.1
Theorem 4.1 embeds n hard single-state transition problems into a layered MDP whose known tree routes the agent to one of n difficult states. Near-optimality under suitable rewards forces the policy to visit each such state substantially.
- Embedded construction: The state space has at most 4n states, and n is chosen as the largest power of two satisfying 4n ≤ S.The lower bound extends to larger state spaces by adding isolated, unvisitable states.
- Embedded construction: The layered construction contains a known dyadic tree leading to n states with unknown transition instances.At each state (x, ℓ0), the learner encounters a transition model P(x) from the single-state hard class.
- Embedded construction: Action 1 moves upward through the tree, whereas actions greater than 1 move rightward from each layered state.
- Policy requirement: ε-suboptimal policies must visit each target state (x, ℓ0) with probability at least 1/2.This follows from the reward structure and holds under the stated ε ≤ 1/4 condition.
- Reduction: The embedding maps a policy for the layered MDP to an n-tuple of policies for the single-state instances.The correspondence is independent of the transition instance and reward vector, though it may depend on ε, n, A, and H.
- Reduction: The proof simulates exploration and planning through the embedded algorithm, transferring correctness to each single-state instance.
D.4 Proof of Lemma D.2
Lemma D.2 is proved by packing many near-uniform transition vectors and reducing accurate reward-free exploration to estimating the hidden packing index for every action. Fano’s inequality then yields the trajectory lower bound.
- Packing construction: The packing uses binary vectors to construct many near-uniform transition distributions indexed by action-specific hidden choices.
- Packing construction: A γ-uncorrelated packing requires distinct vectors to have small pairwise inner products.The definition bounds each distinct pair’s inner product magnitude by 2nγ.
- Information-theoretic lower bound: Fano’s inequality converts the difficulty of estimating the hidden indices into a lower bound on the number of collected trajectories.
- Information-theoretic lower bound: For sufficiently large packing size and suitable parameters, the resulting bound has the required ε^-2 dependence.
- Reduction to estimation: Reward vectors are designed to distinguish the transition vector selected for one action from alternatives associated with other actions.For each action, policies returned for these rewards identify the correct packing index through their initial-action probabilities.
D.4.1 Proof of Lemma D.6
Lemma D.6 constructs the required uncorrelated packing probabilistically. Random binary vectors have concentrated pairwise inner products, and a union bound guarantees a packing of exponential size under the stated parameter condition.
- Probabilistic construction: The construction samples binary vectors randomly and uses concentration bounds to control their pairwise inner products.
- Probabilistic construction: A union bound over all distinct vector pairs establishes that an uncorrelated packing occurs with nonzero probability.
- Packing existence: A γ-uncorrelated packing exists whenever 2 log(M) ≤ nγ^2 − log(4n) − 2 log(A).
- Fano reduction: Fano’s inequality is applied separately to each action by comparing transition laws that differ only when that action is selected.
- Fano reduction: The reduction uses reward vectors whose induced policies reveal the hidden packing index when the packing is sufficiently uncorrelated.