Outcome-Driven Reinforcement Learning via Variational Inference

1 Introduction . Reinforcement learning ( RL ) provides an appealing formalism for automated learning of behavioral skills , but requires considerable care and manual design to use in practice . One particularly delicate decision is the design of the reward function , which has a significant impact on the resulting policy but is largely heuristic in practice , often lacks theoretical grounding , can make effective learning difficult , and may lead to reward mis-specification . To avoid these shortcomings , we propose to circumvent the process of manually specifying a reward function altogether : Instead of framing the reinforcement learning problem as finding a policy that maximizes a heuristically-defined reward function , we express it probabilistically , as inferring a state–action trajectory distribution conditioned on a desired future outcome . By building off of prior work on probabilistic perspectives on RL [ 10 , 23 , 35 , 46 , 47 , 58 ] and goal-directed RL in particular [ 3 , 6 , 19 , 48 ] , we derive a tractable variational objective , an temporal-difference algorithm that provides a shaping-like effect for effective learning , as well as a reward function that captures the semantics of the underlying decision problem and facilitates effective learning . We demonstrate that unlike prior works that proposed inference methods for finding policies that achieve desired outcomes [ 3 , 11 , 19 , 48 ] , the resulting algorithm , Outcome-Driven Actor–Critic ( ODAC ) , is amenable to off-policy learning and applicable to complex , high-dimensional continuous control tasks over finite and infinite horizons . The resulting variational algorithm can be interpreted as an automatic shaping method , where each iteration learns a reward function that automatically provides dense rewards , as we visualize in Figure 1 . In tabular settings , ODAC is guaranteed to converge ∗Equal contribution . † Corresponding authors : tim.rudner @ cs.ox.ac.uk and vitchyr @ berkeley.edu . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . to an optimal policy , and in non-tabular settings with linear Gaussian transition dynamics , the derived optimization objective is convex in the policy , facilitating easier learning . In high-dimensional and non-linear domains , our method can be combined with deep neural network function approximators to yield a deep reinforcement learning method that does not require manual specification of rewards , and leads to good performance on a range of benchmark tasks . Contributions . The core contributions of this paper are the probabilistic formulation of a general framework for inferring policies that lead to desired outcomes and the derivation of a variational objective from which we obtain a novel outcome-driven Bellman backup operator . We show that this Bellman backup operator induces a shaping-like effect which ensures a clear and dense learning signal even in the early stages of training . Crucially , unlike heuristic approaches for incorporating shaping often used in standard RL , this “ shaping ” emerges automatically from variational inference . We demonstrate that the resulting variational objective is a lower bound on the log-marginal likelihood of achieving the outcome given an initial state and that it leads to an off-policy temporal-difference algorithm . We evaluate this algorithm—Outcome-Driven Variational Inference ( ODAC ) —on a range of reinforcement learning tasks without having to manually specify task-specific reward functions . In our experiments , we find that our method results in significantly faster learning across a variety of robot manipulation and locomotion tasks than alternative approaches . 2 Preliminaries . Standard reinforcement learning ( RL ) addresses reward maximization in a Markov decision process ( MDP ) defined by the tuple ( S , A , pS0 , pd , r , γ ) [ 43 , 44 ] , where S and A denote the state and action space , respectively , p0 denotes the initial state distribution , pd is a state transition distribution , r is an immediate reward function , and γ is a discount factor . To sample trajectories , an initial state is sampled according to pS0 , and successive states are sampled from the state transition distribution St+1 ∼ pd ( · | st , at ) and actions from a policy At ∼ π ( · | st ) . We will write T 0 : t = { S0 , A0 , S1 , . . . , St , At } to represent a finite-horizon and T 0 =̇ { At , St+1 } ∞t=0 to represent an infinite-horizon stochastic state–action trajectory , and write τ0 : t = { s0 , a0 , s1 , . . . , st , at } and τ0 =̇ { at , st+1 } ∞t=0 for the respective trajectory realizations . Given a reward function r : S×A → R and discount factor γ ∈ ( 0 , 1 ) , the objective in reinforcement learning is to find a policy π that maximizes the returns , defined as Epπ [ ∑∞ t=0 γ tr ( st , at ) ] , where pπ denotes the distribution of states induced by a policy π. Goal-Conditioned Reinforcement Learning . In goal-conditioned reinforcement learning [ 22 ] , which can be considered a special case of the broader class of stochastic boundary value problems [ 1 , 14 ] , the objective is for an agent to reach some pre-specified goal state , g ∈ S , so that the policy and reward function introduced above become dependent on the goal and are expressed as π ( a | s , g ) and r ( s , a , g ) , respectively . Typically , such a reward function needs to be defined manually , with a common choice being to use a sparse indicator reward r ( s , g ) = I { s = g } . However , this approach presents a number of challenges both in theory and in practice . From a theoretical perspective , the indicator reward will almost surely equal zero for environments with continuous goal spaces and non-trivial stochastic dynamics . From a practical perspective , such sparse rewards can be slow to learn from , as most transitions provide no reward supervision , while manually designing dense reward functions that provide a better learning signal is time-consuming and often based on heuristics . In Section 3 , we will present a framework that addresses these practical and theoretical considerations by casting goal-conditioned RL as probabilistic inference . Q-Learning . Off-policy Q-learning algorithms [ 52 ] allow training policies from data collected under alternate decision rules by estimating the expected return Qπ for a given state–action pair : Qπ ( s , a ) =̇ Epπ [ ∑∞ t=0 γtr ( st , at ) ∣∣∣S0 = s , A0 = a ] . Crucially , the expected return given a state–action pair can be expressed recursively as Qπ ( s , a ) = r ( s , a ) +γ Epπ [ Qπ ( s1 , a1 ) |S0 = s , A0 = a ] , ( 1 ) which makes it possible to estimate the expectation on the right-hand side from single-step transitions . The resulting estimates can then be used to find a policy that results in actions which maximize the expected return Qπ ( s , a ) for all available state–action pairs . 3 Outcome-Driven Reinforcement Learning . In this section , we derive a variational inference objective to infer an approximate posterior policy for achieving desired outcomes . Instead of using the heuristic goal-reaching rewards discussed in Section 2 , we will derive a general framework for inferring actions that lead to desired outcomes by formulating a probabilistic objective , using the tools of variational inference . As we will show in the following sections , we use this formulation to translate the problem of inferring a policy that leads to a desired outcome into a tractable variational optimization problem , which we show corresponds to an RL problem with a well-shaped , dense reward signal from which the agent can learn more easily . We start with a warm-up problem that demonstrates how to frame the task of achieving a desired outcome as an inference problem in a simplified setting . We then describe how to extend this approach to more general settings . Finally , we show that the resulting variational objective can be expressed as a recurrence relation , which allows us to derive an outcome-driven variational Bellman operator and prove an outcome-driven probabilistic policy iteration theorem . 3.1 Warm-up : Achieving a Desired Outcome at a Fixed Time Step . We first consider a simplified problem , where the desired outcome is to reach a specific state g ∈ S at a specific time step t ? when starting from initial state s0 . We can think of the starting state s0 and the desired outcome g as boundary conditions , and the goal is to learn a stochastic policy that induces a trajectory from s0 to g. To derive a control law that solves this stochastic boundary value problem , we frame the problem probabilistically , as inferring a state–action trajectory posterior distribution conditioned on the desired outcome and the initial state . We will show that , by framing the learning problem this way , we obtain an algorithm for learning outcome-driven policies without needing to manually specify a reward function . We consider a model of the state–action trajectory up to and including the desired outcome g , pT̃ 0 : t , St ? ( τ̃0 : t , g | s0 ) =̇ pd ( g | st , at ) p ( at | st ) t−1∏ t′=0 pd ( st′+1 | st′ , at′ ) p ( at′ | st′ ) , where t ? =̇ t+1 , T̃ 0 : t is the state–action trajectory up to an including t but excluding S0 , p ( at | st ) is a conditional action prior , and pd ( st+1 | st , at ) is the environment ’ s state transition distribution . If the dynamics are simple ( e.g. , tabular or Gaussian ) , the posterior over actions can be computed in closed form [ 3 ] , but we would like to be able to infer outcome-driven posterior policies in any environments , including those where exact inference may be intractable . To do so , we start by expressing posterior inference as the variational minimization problem min qT̃ 0 : t|S0 ∈Q̂ DKL ( qT̃ 0 : t|S0 ( · | s0 ) ‖ pT̃ 0 : t|S0 , St ? ( · | s0 , g ) ) , ( 2 ) where DKL ( · ‖ · ) is the KL divergence , and Q̂ denotes the variational family over which to optimize . We consider a family of distributions parameterized by a policy π and defined by qT̃ 0 : t|S0 ( τ̃0 : t | s0 ) =̇π ( at | st ) t−1∏ t′=0 pd ( st′+1 | st′ , at′ ) π ( at′ | st′ ) , ( 3 ) where π ∈ Π , a family of policy distributions , and where ∏t−1 t=0 pd ( st+1 | st , at ) is the true actionconditional state transition distribution up to but excluding the state transition at t ? − 1 , since St ? = g is observed . Under this variational family , the inference problem in Equation ( 2 ) can be equivalently stated as the problem of maximizing the following objective with respect to the policy π : Proposition 1 . Given qT̃ 0 : t|S0 ( τ̃0 : t | s0 ) from Equation ( 3 ) , any state s0 , termination time t ? , and outcome g , solving Equation ( 2 ) is equivalent to maximizing this objective with respect to π ∈ Π : F̄ ( π , s0 , g ) =̇ EqT̃ 0 : t|S0 ( τ̃0 : t | s0 ) [ log pd ( g | st , at ) − t−1∑ t′=0 DKL ( π ( · | st′ ) || p ( · | st′ ) ) ] . ( 4 ) Proof . See Appendix A.1 . A variational problem of this form—which corresponds to finding a posterior distribution over state–action trajectories—can equivalently be viewed as a reinforcement learning problem : Corollary 1 . The objective in Equation ( 4 ) corresponds to KL-regularized reinforcement learning with a time-varying reward function given by r ( st′ , at′ , g , t′ ) =̇ I { t′ = t } log pd ( g | st′ , at′ ) . Corollary 1 illustrates how a reward function emerges automatically from a probabilistic framing of outcome-driven reinforcement learning problems where the sole specification is which variable ( St ? ) should attain which value ( g ) . In particular , Corollary 1 suggests that we ought to learn the environment ’ s state-transition distribution , and view the log-likelihood of achieving the desired outcome given a state–action pair as a “ reward ” that can be used for off-policy learning as described in Section 2 . Importantly—and unlike in model-based RL—such a transition model would not have to be accurate beyond single-step predictions , as it would not be used for planning ( see Appendix B ) . Instead , log pd ( g | st , at ) only needs to be well shaped , which we expect to happen for commonly used model classes . For example , when the dynamics are linear-Gaussian , using a conditional Gaussian model [ 29 ] yields a reward function that is quadratic in St+1 , making the objective convex and thus more amenable to optimization .

Summary: This paper combines the recent work on viewing RL as Probabilistic Inference and the goal-conditioned RL setting. The proposed approach assigns rewards based on the likelihood of reaching the desired state-action pair and propagates this via learning a dynamic discount function based on time elapsed in the episode. The authors propose solving the objective via variational inference and demonstrate this leads to a probabilistic policy evaluation operator. An actor-critic approach is then taken to perform approximate policy iteration.  The authors show the advantage of their methods on high-dimensional continuous control tasks.