SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$f$ -Policy Gradients: A General Framework for Goal Conditioned RL using  $f$ -Divergences
Siddhant Agarwal
The University of Texas at Austin siddhant@cs.utexas.edu
The University of Texas at Austin  
Sony AI  
pstone@cs.utexas.edu
Ishan Durugkar
Sony AI
ishan.durugkar@sony.com
Amy Zhang
The University of Texas at Austin amy.zhang@austin.utexas.edu
Goal-Conditioned Reinforcement Learning (RL) problems often have access to sparse rewards where the agent receives a reward signal only when it has achieved the goal, making policy optimization a difficult problem. Several works augment this sparse reward with a learned dense reward function, but this can lead to sub-optimal policies if the reward is misaligned. Moreover, recent works have demonstrated that effective shaping rewards for a particular problem can depend on the underlying learning algorithm. This paper introduces a novel way to encourage exploration called  $f$ -Policy Gradients, or  $f$ -PG.  $f$ -PG minimizes the f-divergence between the agent's state visitation distribution and the goal, which we show can lead to an optimal policy. We derive gradients for various f-divergences to optimize this objective. Our learning paradigm provides dense learning signals for exploration in sparse reward settings. We further introduce an entropy-regularized policy optimization objective, that we call state-MaxEnt RL (or  $s$ -MaxEnt RL) as a special case of our objective. We show that several metric-based shaping rewards like L2 can be used with  $s$ -MaxEnt RL, providing a common ground to study such metric-based shaping rewards with efficient exploration. We find that  $f$ -PG has better performance compared to standard policy gradient methods on a challenging gridworld as well as the Point Maze and FetchReach environments. More information on our website https://agarwalsiddhant10.github.io/projects/fpg.html.
Reinforcement Learning (RL) algorithms aim to identify the optimal behavior (policy) for solving a task by interacting with the environment. The field of RL has made large strides in recent years (Mnih et al., 2013; Silver et al., 2017; Haarnoja et al., 2018; Ouyang et al., 2022; Wurman et al., 2022) and has been applied to complex tasks ranging from robotics (Gupta et al., 2019), protein synthesis (Jumper et al., 2021), computer architecture (Fawzi et al., 2022) and finance (Liu et al., 2021). Goal-Conditioned RL (GCRL) is a generalized form of the standard RL paradigm for learning a policy that can solve many tasks, as long as each task can be defined by a single rewarding goal state. Common examples of goal-conditioned tasks arise in robotics where the goal states can be a target object configuration for manipulation-based tasks (Kim et al., 2022; Gupta et al., 2019; OpenAI et al., 2021) or a target location for navigation-based tasks (Shah et al., 2020; Gervet et al., 2023).
In any reinforcement learning setup, the task is conveyed to the agent using rewards (Silver et al., 2021). In goal-conditioned RL settings, a common reward function used is 1 when the goal is
achieved and 0 everywhere else. This reward function is sparse and poses a huge learning challenge to obtain the optimal policy without any intermediate learning signal. Prior works (Ng et al., 1999; Ni et al., 2020; Durugkar et al., 2021; Arjona-Medina et al., 2019; Goyal et al., 2019) have augmented the reward function to provide some dense signal for policy optimization. A major issue with augmenting reward functions is that the optimal policy for the new reward function may no longer be optimal under the original, true reward function (Ng et al., 1999). Moreover, it has been shown (Booth et al., 2023) that shaping rewards that improve learning for one learning algorithm may not be optimal for another learning algorithm. Algorithms that learn reward functions (Ni et al., 2020; Durugkar et al., 2021; Zheng et al., 2018) are inefficient because the reward function must first be learned before it can be used for policy optimization. These challenges lead to the following research question: Is there another way to provide dense learning signals for policy optimization other than through dense shaping rewards?
In this work, we look at using divergence minimization between the agent's state visitation and the goal distribution (we assume that each goal can be represented as a distribution, Dirac distribution being the simplest) as an objective to provide additional learning signals. Similar perspectives to policy learning has been explored by prior works (Ziebart et al., 2008; Haarnoja et al., 2017, 2018; Ho & Ermon, 2016; Ni et al., 2020; Ghasemipour et al., 2019; Fu et al., 2017), but they reduce their methods into a reward-centric view. MaxEnt RL methods (Ziebart et al., 2008; Haarnoja et al., 2017, 2018) use the distribution over trajectories rather than state visitations and still suffer from sparsity if the task rewards are sparse. Imitation learning works like those of Ho & Ermon (2016); Fu et al. (2017); Ghasemipour et al. (2019) use a variational lower bound to obtain min-max objectives that require discriminators. These objectives suffer from mathematical instabilities and often require coverage assumptions i.e., abundant overlap between the agent's state visitation distribution and goal distribution. Our method does not rely on discriminators nor does it assume state coverage. It provides dense signals to update the policy even when the agent has not seen the goal. These signals push the policy towards higher entropy state visitations until the goal is discovered.
Our method,  $f$ -PG or  $f$ -Policy Gradient, introduces a novel GCRL framework that aims to minimize a general measure of mismatch (the  $f$ -divergence) between the agent's state visitation distribution and the goal distribution. We prove that minimizing the  $f$ -divergence (for some divergences) recovers the optimal policy. The analytical gradient for the objective looks very similar to a policy gradient which allows us to leverage established methods from the policy gradient literature to come up with an efficient algorithm for goal-conditioned RL. We show the connection of our method to the commonly used metric-based shaping rewards for GCRL like L2 rewards. We show that a special case of  $f$ -PG jointly optimizes for maximization of a reward and the entropy of the state-visitation distribution thus introducing state-MaxEnt RL (or s-MaxEnt RL). Using a sparse gridworld, we establish the benefits of using  $f$ -PG as a dense signal to explore when the agent has not seen the goal. We also demonstrate that our framework can be extended to continuous state spaces and scale to larger and higher-dimensional state spaces in maze navigation and manipulation tasks.
Our key contributions are 1) developing a novel algorithm for goal-conditioned RL that provably produces the optimal policy, 2) connecting our framework to commonly known metric-based shaping rewards, 3) providing a new perspective to RL (s-MaxEnt RL) that focuses on maximizing the entropy of the state-visitation distribution and 4) empirical evidence demonstrating its ability to provide dense learning signals and scale to larger domains.
This section goes over the standard goal-conditioned reinforcement learning formulation and the f-divergences that will be used in the rest of the paper.
Goal-conditioned reinforcement learning. This paper considers an agent in a goal-conditioned MDP (Puterman, 1990; Kaelbling, 1993). A goal-conditioned MDP is defined as a tuple  $\langle S, \mathcal{G}, \mathcal{A}, P, r, \gamma, \mu_0, \rho_g \rangle$  where  $S$  is the state space,  $\mathcal{A}$  is the action space,  $P: S \times \mathcal{A} \longmapsto \Delta(S)$  is the transition probability ( $\Delta(\cdot)$  denotes a probability distribution over a set),  $\gamma \in [0,1)$  is the discount factor,  $\mu_0$  is the distribution over initial states,  $\mathcal{G} \subset S$  is the set of goals, and  $\rho_g: \Delta(\mathcal{G})$  is the distribution over goals. At the beginning of an episode, the initial state  $s_0$  and the goal  $g$  are sampled from the distributions  $\mu_0$  and  $\rho_g$ . The rewards  $r: S \times \mathcal{G} \longmapsto \mathbb{R}$  are based on the state the agent visits and conditioned on the goal specified during that episode. This work focuses on sparse rewards, where  $r(s', g) = 1$  when  $s' = g$ , and is  $r(s', g) = 0$  otherwise. In continuous domains, the equality is relaxed to  $s' \in \mathcal{B}(g, r)$  where  $\mathcal{B}(g, r)$  represents a ball around the goal  $g$  with radius  $r$ .
A trajectory  $\tau$  is defined as the sequence  $(s_0, a_0, s_1, \ldots, s_{T-1}, a_{T-1}, s_T)$ . The return  $H_g(s)$  is defined as the cumulative undiscounted rewards  $H_g(s) := \sum_{t=0}^{T} [r(s_{t+1}, g) | s_0 = s]$ , where  $T$  is the length of a trajectory. We will assume the trajectory ends when a maximum number of policy steps  $(T)$  have been executed. The agent aims to learn a policy  $\pi: \mathcal{S} \times \mathcal{G} \longmapsto \Delta(\mathcal{A})$  that maximises the expected return  $\mathbb{E}_{\pi, s_0}[H_g(s_0)]$ . The optimal policy  $\pi^* = \arg \max_{\pi_\theta \in \Pi} \mathbb{E}_{\pi, s_0}[H_g(s_0)]$ , where the space of policies  $\Pi$  is defined by a set of parameters  $\theta \in \Theta$ .
Distribution matching approach to goal-conditioned RL. The distribution over goal-conditioned trajectories is defined as  $p_{\theta}(\tau; g) = p(s_0)\Pi_{t=0}^{T}p(st|s_{t-1}, a_{t-1})\pi_{\theta}(a_t|s_t; g)$ . The trajectory-dependent state visitation distribution is defined as  $\eta_{\tau}(s)$ . It is the number of times the state  $s$  is visited in the trajectory  $\tau$ . The agent's goal-conditioned state visitation can then be defined as:
$$
\begin{array}{l} p _ {\theta} (s; g) = \frac {\int p _ {\theta} (\tau ; g) \eta_ {\tau} (s) d \tau}{Z} (1) \\ = \frac {\int \Pi p \left(s _ {t + 1} \mid s _ {t} , a _ {t}\right) \pi_ {\theta} \left(a _ {t} \mid s _ {t} ; g\right) \eta_ {\tau} (s)}{\int \int \Pi p \left(s _ {t + 1} \mid s _ {t} , a _ {t}\right) \pi_ {\theta} \left(a _ {t} \mid s _ {t} ; g\right) \eta_ {\tau} (s) d \tau d s} d \tau . (2) \\ \end{array}
$$
The goal  $g$  defines an idealized target distribution  $p_{g} : \Delta(S)$ , considered here as a Dirac distribution which places all the probability mass at the goal state  $p_{g} = \delta(g)$ . Such a formulation has been used previously in approaches to learn goal-conditioned policies (Durugkar et al., 2021). This work focuses on minimizing the mismatch of an agent's goal-conditioned state visitation distribution  $p_{\theta}(s; g)$  to this target distribution  $p_{g}$ . In this paper, we will be using  $p_{\theta}$  and  $p_{\pi}$  interchangeably i.e.,  $p_{\theta}$  corresponds to the visitation distribution induced by policy  $\pi$  that is parameterized by  $\theta$ .
To do so, this paper considers a family of methods that compare the state-visitation distribution induced by a goal-conditioned policy and the ideal target distribution for that goal  $g$ , called  $f$ -divergences.  $f$ -divergences are defined as (Polyanskiy & Wu, 2022),
$$
D _ {f} (P | | Q) = \int_ {P > 0} P (x) f \left(\frac {Q (x)}{P (x)}\right) d x - f ^ {\prime} (\infty) Q ([ P (x) = 0 ]), \tag {3}
$$
where  $f$  is a convex function with  $f(1) = 0$ .  $f^{\prime}(\infty)$  is not defined (is  $\infty$ ) for several  $f$ -divergences and so it is a common assumption that  $Q = 0$  wherever  $P = 0$ . Table 1 shows a list of commonly used  $f$ -divergences with corresponding  $f$  and  $f^{\prime}(\infty)$ .
f-divergence	Df(P||Q)	f(u)	f'(u)	f'(∞)
FKL	∫P(x) log P(x)/Q(x)dx	u log u	1 + log u	Undefined
RKL	∫Q(x) log Q(x)/P(x)dx	- log u	- 1/u	0
JS	1/2 ∫P(x) log 2P(x)/P(x)+Q(x)+Q(x)/P(x)+Q(x)dx + Q(x) log 2Q(x)/P(x)+Q(x)dx	u log u-(1 + u) log 1+u/2	log 2u/1+u	log 2
χ²	1/2 ∫Q(x)(P(x)/Q(x) - 1)2dx	1/2(u - 1)2	u	Undefined
Table 1: Selected list of  $f$ -divergences  $D_{f}(P||Q)$  with generator functions  $f$  and their derivatives  $f'$ , where  $f$  is convex, lower-semicontinuous and  $f(1) = 0$ .
Shaping Rewards. Our work is related to a separate class of techniques that augment the sparse reward function with dense signals. Ng et al. (1999) proposes a way to augment reward functions without changing the optimal behavior. Intrinsic Motivation (Durugkar et al., 2021; Bellemare et al., 2016; Singh et al., 2010; Barto, 2013) has been an active research area for providing shaping rewards. Some work (Niekum, 2010; Zheng et al., 2018) learn intrinsic or alternate reward functions for the underlying task that aim to improve agent learning performance while others (Durugkar et al., 2021; Ni et al., 2020; Goyal et al., 2019) learn augmented rewards based on distribution matching. AIM (Durugkar et al., 2021) learns a potential-based shaping reward to capture the time-step distance but requires a restrictive assumption about state coverage, especially around the goal while we do not make any such assumption. Recursive classification methods (Eysenbach et al., 2021, 2020) use future state densities as rewards. However, these methods will fail when the agent has never seen the goal. Moreover, in most of these works, the reward is not stationary (is dependent on the policy)
which can lead to instabilities during policy optimization. GoFAR (Ma et al., 2022) is an offline goal-conditioned RL algorithm that minimizes a lower bound to the KL divergence between  $p_{\theta}(s)$  and the  $p_g(s)$ . It computes rewards using a discriminator and uses the dual formulation utilized by the DICE family (Nachum et al., 2019), but reduces to GAIL (Ho & Ermon, 2016) in the online setting, requiring coverage assumptions. Our work also minimizes the divergence between the agent's visitation distribution and the goal distribution, but we provide a new formulation for on-policy goal-conditioned RL that does not require a discriminator or the same coverage assumptions.
Policy Learning through State Matching. We first focus on imitation learning where the expert distribution  $p_E(s, a)$  is directly inferred from the expert data. GAIL (Ho & Ermon, 2016) showed that the inverse RL objective is the dual of state-matching. f-MAX (Ghasemipour et al., 2019) uses f-divergence as a metric to match the agent's state-action visitation distribution  $p_{\pi}(s, a)$  and  $p_E(s, a)$ . Ke et al. (2019); Ghasemipour et al. (2019) shows how several commonly used imitation learning methods can be reduced to a divergence minimization. But all of these methods optimize a lower bound of the divergence which is essentially a min-max bilevel optimization objective. They break the min-max into two parts, fitting the density model to obtain a reward that can be used for policy optimization. But these rewards depend on the policy, and should not be used by RL algorithms that assume stationary rewards. f-IRL (Ni et al., 2020) escapes the min-max objective but learns a reward function that can be used for policy optimization. We do not aim to learn a reward function but rather directly optimize for a policy using dense signals from an  $f$ -divergence objective.
In reinforcement learning, the connections between entropy regularized MaxEnt RL and the minimization of reverse KL between agent's trajectory distribution,  $p_{\pi}(\tau)$ , and the "optimal" trajectory distribution,  $p^{*}(\tau) \propto e^{r(\tau)}$  has been extensively studied Ziebart (2010); Ziebart et al. (2008); Kappen et al. (2012); Levine (2018); Haarnoja et al. (2018). MaxEnt RL optimizes for a policy with maximum entropy but such a policy does not guarantee maximum coverage of the state space. Hazan et al. (2018) discusses an objective for maximum exploration that focuses on maximizing the entropy of the state-visitation distribution or KL divergence between the state-visitation distribution and a uniform distribution. A few works like Durugkar et al. (2023, 2021); Ma et al. (2022), that have explored state-matching for reinforcement learning, have been discussed above. Several works like (Belousov & Peters, 2018, 2019; Touati et al., 2020) have used divergence to constraint the policy improvement steps making the updates more stable.
Limitations of Markov Rewards. Our work looks beyond the maximization of a Markov reward for policy optimization. The learning signals that we use are non-stationary. We thus discuss the limitations of using Markov rewards for obtaining the optimal policy. There have been works (Abel et al., 2021; Clark & Amodei, 2016; Icarte et al., 2018, 2021) that express the difficulty in using Markov rewards. Abel et al. (2021) proves that there always exist environment-task pairs that cannot be described using Markov rewards. Reward Machines (Icarte et al., 2018) create finite automata to specify reward functions and can specify Non-Markov rewards as well but these are hand-crafted.
4  $f$ -Policy Gradient
In this paper, we derive an algorithm where the agents learn by minimizing the following  $f$ -divergence:
$$
J (\theta) = D _ {f} \left(p _ {\theta} (s) \mid \mid p _ {g} (s)\right) \tag {4}
$$
In this section, we shall derive an algorithm to minimize  $J(\theta)$  and analyze the objective more closely in the subsequent section. Unlike f-max (Ghasemipour et al., 2019), we directly optimize  $J(\theta)$ . We differentiate  $J(\theta)$  with respect to  $\theta$  to get this gradient. We follow a similar technique as Ni et al. (2020) to obtain the analytical gradient of  $J(\theta)$ .
Theorem 4.1. The gradient of  $J(\theta)$  as defined in Equation 4 is given by,
$$
\nabla_ {\theta} J (\theta) = \mathbb {E} _ {\tau \sim p _ {\theta} (\tau)} \left[ \left[ \sum_ {t = 1} ^ {T} \nabla_ {\theta} \log \pi_ {\theta} \left(a _ {t} \mid s _ {t}\right) \right] \left[ \sum_ {t = 1} ^ {T} f ^ {\prime} \left(\frac {p _ {\theta} \left(s _ {t}\right)}{p _ {g} \left(s _ {t}\right)}\right) \right] \right]. \tag {5}
$$
The gradient looks exactly like policy gradient with rewards  $-f'\left(\frac{p_{\theta}(s_t)}{p_g(s_t)}\right)$ . However, this does not mean that we are maximizing  $J^{RL}(\theta) = \mathbb{E}_{\tau \sim p_{\theta}(\tau)}\left[-f'\left(\frac{p_{\theta}(s_t)}{p_g(s_t)}\right)\right]$ . This is because the gradient of  $J^{RL}(\theta)$  is not the same as  $\nabla_{\theta}J(\theta)$ . For Dirac goal distributions, the gradient in Equation 5 cannot be
used (as  $f' \left( \frac{p_{\theta}(s_t)}{p_g(s_t)} \right)$  will not be defined when  $p_g(s_t) = 0$ ). We can use the definition of  $f$ -divergence in Equation 3 to derive a gradient for such distributions.
The gradient is obtained in terms of the state visitation frequencies  $\eta_{\tau}(s)$ . Further examination of the gradient leads to the following theorem,
Theorem 4.2. Updating the policy using the gradient (Equation 5) maximizes  $\mathbb{E}_{p\theta}[\eta_{\tau}(g)]$
Theorem 4.2 provides another perspective for  $f$ -Policy Gradient -  $\eta_{\tau}(g)$  is equivalent to the expected return for a goal-based sparse reward, hence optimizing the true goal-conditioned RL objective. We shall prove the optimality of the policy obtained from minimizing  $J(\theta)$  in the next section.
In practice, a Dirac goal distribution can be approximated by clipping off the zero probabilities at  $\epsilon$ , similar to Laplace correction. Doing so, we will be able to use dense signals from the gradient in Equation 5 while still producing the optimal policy. This approximation is different from simply adding an  $\epsilon$  reward at every state. This is because the gradients are still weighed by  $f'\left(\frac{p_{\theta}(s_t)}{\epsilon}\right)$  which depends on  $p_{\theta}(s_t)$ .
Simply optimizing  $J(\theta)$  is difficult because it faces similar issues to REINFORCE (Williams & Peng, 1991). A major shortcoming of the above gradient computation is that it requires completely on-policy updates. This requirement will make learning sample inefficient, especially when dealing with any complex environments. However, there have been a number of improvements to naïve policy gradients that can be used. One approach is to use importance sampling (Precup, 2000), allowing samples collected from a previous policy  $\pi_{\theta'}$  to be used for learning. To reap the benefits of importance sampling, we need the previous state-visitation distributions to compute  $f' \left( \frac{p_{\theta}(s)}{p_g(s)} \right)$ . Hence, we need to ensure that the current policy does not diverge much from the previous policy. This condition is ensured by constraining the KL divergence between the current policy and the previous policy. We use the clipped objective similar to Proximal Policy Optimization (Schulman et al., 2017), which has been shown to work well with policy gradients. PPO has shown that the clipped loss works well even without an explicit KL constraint in the objective. The gradient used in practice is,
$$
\nabla_ {\theta} J (\theta) = \mathbb {E} _ {s _ {t}, a _ {t} \sim p _ {\theta^ {\prime}} (s _ {t}, a _ {t})} \left[ \min  \left(r _ {\theta} \left(s _ {t}\right) F _ {\theta^ {\prime}} \left(s _ {t}\right), c l i p \left(r _ {\theta} \left(s _ {t}\right), 1 - \epsilon , 1 + \epsilon\right) F _ {\theta^ {\prime}} \left(s _ {t}\right)\right) \right] \tag {6}
$$
where  $r_{\theta}(t) = \frac{\pi_{\theta}(a_t|s_t)}{\pi_{\theta'}(a_t|s_t)}$  and  $F_{\theta '}(s_t) = \sum_{t' = t}^{T}\gamma^{t'}f'\left(\frac{p_{\theta'}(s_t)}{p_g(s_t)}\right)$ . The derivation for this objective is provided in Appendix B.  $\gamma$  is added to improve the stability of gradients and to prevent the sum of  $f^{\prime}\left(\frac{p_{\theta^{\prime}}(s_t)}{p_g(s_t)}\right)$  from exploding.
For the purpose of this paper, we use kernel density estimators to estimate the goal distribution and the agent's state visitation distribution. We may also use discriminators to estimate the ratio of these densities like Ho & Ermon (2016); Fu et al. (2017); Ghasemipour et al. (2019). But unlike these methods, we will not be incorrectly breaking a minmax objective. In our case, the estimate of the gradient requires the value of the ratio of the two distributions and does not make any assumptions about the stationarity of these values. While the adversarial methods break the minmax objective and assume the discriminator
to be fixed (and rewards stationary) during policy optimization.
Algorithm 1  $f$  -PG
Let,  $\pi_{\theta}$  be the policy,  $G$  be the set of goals,  $B$  be a buffer  
for  $i = 1$  to num_iter do  
 $B \gets []$   
for  $j = 1$  to num_traj_per_iter do  
Sample  $g$ , set  $p_g(s)$   
Collect goal conditioned trajectories,  $\tau : g$   
Fit  $p_{\theta}(s)$  using KDE on  $\tau$   
Store  $f' \left( \frac{p_{\theta}(s)}{p_g(s)} \right)$  for each  $s$  in  $\tau$ $B \gets B + \{\tau : g\}$   
end for  
for  $j = 1$  to num_policy Updates do  
 $\theta \gets \theta - \alpha \nabla_{\theta} J(\theta)$  (Equation 6)  
end for
5 Theoretical analysis of  $f$ -PG
In this section, we will first show that minimizing the f-divergence between the agent's state visitation distribution and goal distribution yields the optimal policy. We will further analyze the connections to metric based shaping rewards and implicit exploration boost from the learning signals. For the rest
of the paper, we will refer to  $f$ -PG using FKL divergence as  $fkl$ -PG,  $f$ -PG using RKL divergence as  $rkl$ -PG and so on.
5.1 Analysis of  $J(\theta)$
This section shows that the policy obtained by minimizing an  $f$ -divergence between the agent's state visitation distribution and the goal distribution is the optimal policy.
Theorem 5.1. The policy that minimizes  $D_{f}(p_{\pi}||p_{g})$  for a convex function  $f$  with  $f(1) = 0$  and  $f'(\infty)$  being defined, is the optimal policy.