Title: f-Policy Gradients: A General Framework for Goal-Conditioned RL using f-Divergences
Venue: NeurIPS 2023
Accepted: yes
Keywords: Goal Conditioned Reinforcement Learning

Shaping Rewards

Reward Design
Abstract: 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.

f -Policy Gradients: A General Framework for Goal Conditioned RL using f -Divergences

Abstract

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.

1 Introduction

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.

2 Background

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 ⟨S , G , A , P, r, γ, µ 0 , ρ g ⟩ where S is the state space, A is the action space, P : S × A ↦- → ∆( S ) is the transition probability ( ∆( · ) denotes a probability distribution over a set), γ ∈ [0 , 1) is the discount factor, µ 0 is the distribution over initial states, G ⊂ S is the set of goals, and ρ g : ∆( 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 µ 0 and ρ g . The rewards r : S × G ↦- → 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 ′ ∈ B ( g, r ) where B ( g, r ) represents a ball around the goal g with radius r .

A trajectory τ is defined as the sequence ( s 0 , a 0 , s 1 , . . . , s T -1 , a T -1 , s T ) . The return H g ( s ) is defined as the cumulative undiscounted rewards H g ( s ) := ∑ T t =0 [ 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 π : S × G ↦- → ∆( A ) that maximises the expected return E π,s 0 [ H g ( s 0 )] . The optimal policy π ∗ = arg max π θ ∈ Π E π,s 0 [ H g ( s 0 )] , where the space of policies Π is defined by a set of parameters θ ∈ Θ .

Distribution matching approach to goal-conditioned RL. The distribution over goal-conditioned trajectories is defined as p θ ( τ ; g ) = p ( s 0 )Π T t =0 p ( s t | s t -1 , a t -1 ) π θ ( a t | s t ; g ) . The trajectorydependent state visitation distribution is defined as η τ ( s ) . It is the number of times the state s is visited in the trajectory τ . The agent's goal-conditioned state visitation can then be defined as:

The goal g defines an idealized target distribution p g : ∆( S ) , considered here as a Dirac distribution which places all the probability mass at the goal state p g = δ ( 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 θ ( s ; g ) to this target distribution p g . In this paper, we will be using p θ and p π interchangeably i.e., p θ corresponds to the visitation distribution induced by policy π that is parameterized by θ .

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),

where f is a convex function with f (1) = 0 . f ′ ( ∞ ) is not defined (is ∞ ) 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 ′ ( ∞ ) .

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 .

3 Related Work

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 θ ( 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 π ( 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 π ( τ ) , and the 'optimal" trajectory distribution, p ∗ ( τ ) ∝ e r ( τ ) 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:

In this section, we shall derive an algorithm to minimize J ( θ ) and analyze the objective more closely in the subsequent section. Unlike f-max (Ghasemipour et al., 2019), we directly optimize J ( θ ) . We differentiate J ( θ ) with respect to θ to get this gradient. We follow a similar technique as Ni et al. (2020) to obtain the analytical gradient of J ( θ ) .

Theorem 4.1. The gradient of J ( θ ) as defined in Equation 4 is given by,

The gradient looks exactly like policy gradient with rewards -f ′ ( p θ ( s t ) p g ( s t ) ) . However, this does not mean that we are maximizing J RL ( θ ) = ❊ τ ∼ p θ ( τ ) [ -f ′ ( p θ ( s t ) p g ( s t ) )] . This is because the gradient of J RL ( θ ) is not the same as ∇ θ J ( θ ) . For Dirac goal distributions, the gradient in Equation 5 cannot be

used (as f ′ ( p θ ( s t ) p g ( s t ) ) 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 η τ ( s ) . Further examination of the gradient leads to the following theorem,

Theorem 4.2. Updating the policy using the gradient (Equation 5) maximizes ❊ p θ [ η τ ( g )] . Theorem 4.2 provides another perspective for f -Policy Gradient η τ ( 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 ( θ ) in the next section.

In practice, a Dirac goal distribution can be approximated by clipping off the zero probabilities at ϵ , 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 ϵ reward at every state. This is because the gradients are still weighed by f ′ ( p θ ( s t ) ϵ ) which depends on p θ ( s t ) .

Simply optimizing J ( θ ) 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 π ′ to be used for learning. To reap the benefits

θ of importance sampling, we need the previous state-visitation distributions to compute f ′ ( p θ ( s ) p g ( s ) ) .

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,

where r θ ( t ) = π θ ( a t | s t ) π θ ′ ( a t | s t ) and F θ ′ ( s t ) = ∑ T t ′ = t γ t ′ f ′ ( p θ ′ ( s t ) p g ( s t ) ) . The derivation for this objective is provided in Appendix B. γ is added to improve the stability of gradients and to prevent the sum of f ′ ( p θ ′ ( s t ) p g ( s t ) ) 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

Algorithm 1 f -PG

to be fixed (and rewards stationary) during policy optimization.

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 ( θ )

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 π || p g ) for a convex function f with f (1) = 0 and f ′ ( ∞ ) being defined, is the optimal policy.

The proof for Theorem 5.1 is provided in Appendix A. The Theorem states that the policy obtained by minimizing the f -divergence between the agent's state-visitation distribution and the goal distribution is the optimal policy for a class of convex functions defining the f -divergence with f ′ ( ∞ ) defined. It thus makes sense to minimize the f -divergence between the agent's visitation and the goal distribution. It must be noted that the objective does not involve maximizing a reward function. Note that the condition that f ′ ( ∞ ) is defined is not true for all f -divergences. The common f -divergences like RKL, TV, and JS have f ′ ( ∞ ) defined rkl -PG, tv -PG, and js -PG will produce the optimal policy.

Forward KL divergence (FKL) has f = u log u and so does not have f ′ ( ∞ ) defined. Does this mean that the policy obtained by minimizing the FKL divergence is not optimal? Lemma 5.1 (proof in Appendix A) shows that the policy obtained maximizes the entropy of the agent's state-visitation distribution along with maximizing a reward of log p g ( s ) .

Lemma 5.1. fkl -PG produces a policy that maximizes the reward log p g ( s ) along with the entropy of the state-visitation distribution.

A similar result can be shown for χ 2 -divergence as well. It must be understood that Lemma 5.1 does not mean that fkl -PG is the same as the commonly studied MaxEnt RL.

Differences from MaxEnt RL : MaxEnt RL, as studied in Haarnoja et al. (2017, 2018), maximizes the entropy of the policy along with the task reward to achieve better exploration. However, maximizing the entropy of the policy does not imply maximum exploration. Hazan et al. (2018) shows that maximizing the entropy of the state-visitation distribution provably provides maximum exploration. Lemma 5.1 shows that fkl -PG maximizes the entropy of the state-visitation distribution along with the reward making it better suited for exploration. To distinguish our work, we call the MaxEnt RL, as discussed in works like Haarnoja et al. (2017, 2018), as π -MaxEnt RL because it only focuses on the entropy of the policy. On the other hand, fkl -PG maximizes the entropy of the state-visitation distribution so we call it state -MaxEnt RL or s -MaxEnt RL . Similarly, sa -MaxEnt RL can be defined to maximize the entropy of the state-action visitation distribution.

(a)

s

(b)

π

Since the agent's state visitation distribution depends on both the policy and the dynamics, simply increasing the entropy of the policy (without considering the dynamics) will not ensure that the agent will visit most of the states or will have a state-visitation distribution with high entropy. In Figure 1, we compare the efficiencies of π -MaxEnt RL and s -MaxEnt RL to explore around a wall in a discrete gridworld. The initial and the goal distributions ( highlighted in green and red respectively)

P

are separated by a wall. This environment is further discussed in Section 6.1 and Appendix C. Figure 1 shows the evolution of the agent's state-visitation distribution with training for s -MaxEnt RL ( fkl -PG) and π -MaxEnt RL (Soft Q Learning (Haarnoja et al., 2017))

Metric-based Shaping Reward: A deeper look into Lemma 5.1 shows that an appropriate choice of p g ( s ) can lead to entropy maximizing policy optimization with metric-based shaping rewards. Define the goal distribution as p g ( s ) = e f ( s ; g ) where f ( s ; g ) captures the metric of the underlying space. Then the fkl -PG objective becomes, min D FKL ( p θ , p g ) = max ❊ p θ [ f ( s ; g )] -❊ p θ [log p θ ] . (7) The above objective maximizes the reward f ( s ; g ) along with the entropy of the agent's state visitation distribution. For an L2 Euclidean metric, f ( s ; g ) will be -|| s -g || 2 2 which is the L2 shaping reward, and the goal distribution will be Gaussian. If the goal distribution is Laplacian, the corresponding shaping reward will be the L1 norm.

AIM (Durugkar et al., 2021) used a potential-based shaping reward based on a time step quasimetric. If we define f ( s ; g ) as a Lipschitz function for the time step metric maximizing at s = g , we end up optimizing for the AIM reward along with maximizing the entropy of the state-visitation distribution.

5.2 Analysis of the learning signals

f -PG involves a learning signal f ′ ( p θ ( s ) p g ( s ) ) to weigh the gradients of log probabilities of the policy. Since we are minimizing the objective (in contrast to policy gradients) the visitation will be pushed towards states with lower values of the learning signal. It is thus important to understand how f ′ ( p θ ( s ) p g ( s ) ) behaves for goal-conditioned RL settings. During the initial stages of training, the agent visits regions with very low p g . For such states, the signal has a higher value compared to the states that have lower p θ , i.e., the unexplored states. This is because for any convex function f , f ′ ( x ) is an increasing function, so minimizing f ′ ( p θ ( s ) p g ( s ) ) will imply minimizing p θ ( s ) for the states with low p g ( s ) . The only way to do this is to increase the entropy of the state-visitation distribution, directly making the agent explore new states. As long as there is no significant overlap between the two distributions, it will push p θ down to a flatter distribution until there is enough overlap with the goal distribution when it will pull back the agent's visitation again to be closer to the goal distribution.

This learning signal should not be confused with reward in reinforcement learning. It is non-stationary and non-Markovian as it depends on the policy. More importantly, we are not maximizing this signal, just using it to weigh the gradients of the policy.

In the following example, we shall use the Reacher environment (Todorov et al., 2012) to illustrate how our learning signal ( f ′ ( p θ ( s ) p g ( s ) )) varies as the agent learns. We will also show how this signal can push for exploration when the agent has not seen the goal yet. Consider the We fix the goal at (-0.21, 0) and show how the learning signal evolves with the policy. While Figure 2 shows the evolution of -f ′ ( p θ ( s ) p g ( s ) ) (note the negation) for fkl -PG, the rest can be found in Appendix D.

The value of f ′ ( p θ ( s ) p g ( s ) ) is highest where the agent's visitation is high and lower where the agent is not visiting. f ′ ( p θ ( s ) p g ( s ) ) has the lowest value at the goal. As the policy converges to the optimal policy, the

regions where the state-visitation distribution is considerably high (towards the bottom-right in the figure), the value for f ′ ( p θ ( s ) p g ( s ) ) decreases for those states (to still push for exploration) but its value at the goal is low enough for the policy to converge.

6 Experiments

Our experiments evaluate our new framework ( f -PG) as an alternative to conventional reward maximization for goal-conditional RL. We pose the following questions:

Does f -PG provide sufficient signals to explore in otherwise challenging sparse reward settings?

How well does our framework perform compared to discriminator-based approaches?

Can our framework scale to larger domains with continuous state spaces and randomly generated goals?

How do different f -divergences affect learning?

The first two questions are answered using a toy gridworld environment. The gridworld has a goal contained in a room which poses a significant exploration challenge. We also show how the dense signal to the gradients of the policy evolves during training on a continuous domain like Reacher. To answer the third question, our framework is compared with several baselines on a 2D Maze solving task (Point Maze). Additionally, we scale to more complex tasks such as FetchReach Plappert et al. (2018) and an exploration-heavy PointMaze.

6.1 Gridworld

We use a gridworld environment to compare and visualize the effects of using different shaping rewards for exploration. We discussed this environment briefly in Section 5.1. The task is for the agent to reach the goal contained in a room. The only way to reach the goal is to go around the wall. The task reward is 1 when the agent reaches the room otherwise it is 0 . The state space is simply the ( x, y ) coordinates of the grid and the goal is fixed. A detailed description of the task is provided in Appendix C. Although the environment seems simple, exploration here is very difficult as there is no incentive for the agent to go around the wall.

(a)

fkl

-PG

(b)

rkl

-PG

(c) AIM

9

target

9

target

6

target

9

target

8

agent

8

agent

8

agent

8

agent

7

7

7

7

6

6

6

6

5

5

5

5

4

4

4

4

3

0

3

3

3

2

2

2

2

1

O

1

1

1

0

0

0

0

0

1

2

3

4

5

6

7

8

6

0

1

2

3

4

5

6

7

8

6

1

2

3

4

5

6

7

8

6

0

1

2

3

4

5

6

7

8

6

Our framework is compared against AIM (Durugkar et al., 2021), which initially introduced this environment and uses a shaping reward obtained from state-matching to solve it, and GAIL (Ho & Ermon, 2016), which uses a discriminator to learn the probability of a state being the goal state. We provide a comparison to other recent methods in Appendix C. All the baselines are implemented on top of Soft Q Learning (Haarnoja et al., 2017) which along with maximizing the augmented rewards, also maximizes the entropy of the policy while f -PG is implemented as an on-policy algorithm without any extrinsic entropy maximization objective. It can be seen from Figure 3 that, f -PG can explore enough to find the way around the room which is difficult for methods like GAIL even after the entropy boost. AIM learns a potential function and can also find its way across the wall. As expected, fkl -PG converges to the policy maximizing the entropy of the state visitation while rkl -PG produces the optimal state visitation as expected from Theorem 5.1. This simple experiment clearly illustrates two things: (1) f -PG can generate dense signals to explore the state space and search for the goal and (2) although discriminator-based methods like GAIL try to perform state-matching, they fail to explore the space well.

6.2 Point Maze

While the gridworld poses an exploration challenge, the environment is simple and has only one goal. This experiment shows that f -PG scales to larger domains with continuous state space and a large set of goals. We use the Point Maze environments (Fu et al., 2020) which are a set of offline RL environments, and modify it to support our online algorithms. The state space is continuous and consists of the position and velocity of the agent and the goal. The action is the force applied in each direction. There are three variations of the environment namely PointMazeU , PointMazeMedium , PointMazeLarge . For the details of the three environments, please refer to Appendix E.

We compare f -PG with several goal-based shaping reward, (used alongside the task reward as described in Ng et al. (1999)) to optimize a PPO policy 1 . The rewards tried (along with their abbreviations in the plots) are AIM (Durugkar et al., 2021)( aim ), GAIL (Ho & Ermon, 2016)( gail ), AIRL (Fu et al., 2017)( airl ) and F-AIRL (Ghasemipour et al., 2019)( fairl ). All these methods employ a state-matching objective. AIM uses Wasserstein's distance while the rest use some form of f -divergence. But, all of them rely on discriminators. Along with these baselines, we experiment using our learning signal as a shaping reward ( fkl-rew ). Additionally, we also compare with PPO being optimized by only the task reward ( none ). For our method, we have only shown results for fkl -PG. For the rest of the possible f -divergences, refer to Section 6.4.

PointMazeU

PointMazeMedium

PointMazeLarge

1.0

0.8

0.6

fkl-pg

Rate

0.8

Rate

0.6

Rate

0.5

0.4

l!e6

airl

Success

0.6

Success

0.4

Success

0.3

aim

none

0.4

0.2

0.2

fairl

0.1

0.2

0.0

0

2

4

6

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.5

1.0

1.5

Timesteps

1e6

Timesteps

1e7

Timesteps

1e7

Figure 4 (plotting mean and std-dev for 3 seeds) clearly illustrates that fkl -PG is able to perform well in all three environments. In fact, it performs better than the baselines in the more difficult environments. It can also be seen that shaping rewards can often lead to suboptimal performance as none is higher than a few of the shaping rewards. As expected, the curve fkl-new performs poorly. In the simpler PointMazeU environment, the performance for most of the shaping rewards are similar (along with none ) but in more complex PointMazeMedium and PointMazeLarge, a lot of these shaping rewards fail.

6.3 Scaling to Complex Tasks

We scale our method to more complex tasks such as FetchReach (Plappert et al., 2018) and a difficult version of PointMaze. In the PointMaze environments used in the previous section, distributions from which the initial state and the goal are sampled, have a significant overlap easing the exploration. We modify these environments to ensure a significant distance between the sampled goal distributions and the agent's state-visitation distribution as shown in Figure 5 (top), making exploration highly challenging. Figure 5 (bottom) shows the comparison of fkl -PG with GAIL (Ho & Ermon, 2016) and AIM (Durugkar et al., 2021).

The following can be concluded from these experiments: (1) The discriminative-based methods heavily depend on coverage assumptions and fail in situations where there is no significant overlap between the goal distribution and the agent's state visitation distribution. fkl -PG does not depend on any such assumptions. (2) f -PG is considerably more stable than these baselines (as indicated by the variance of these methods).

6.4 Comparing different f -divergences

We perform an ablation to compare different f -divergences on their performances on the three Point Maze environments. Figure 6 (plotting mean and std-dev for 3 seeds) show that, empirically, fkl -PG performs the best followed by χ 2 -PG. Interestingly, both of these do not guarantee optimal policies

3

PointMazeLargeTough

PointMazeMediumTough

FetchReach-vo

1.0

0.8

1.0

0.8

Success Rate

0.6

SuccessRate

Success Rate

0.8

0.6

0.4

0.4

0.6

0.4

0.2

AIM

0.2

0.2

FPG

F00

0.0

0.0

0.5

1.0

1.5

2.0

0.0

0.5

1.0

1.5

2.0

0

100000

200000

300000

Timesteps

1e7

Timesteps

1e7

Timesteps

PointMazeU

PointMazeMedium

PointMazeLarge

1.0

fkl-pg

0.6

rkl-pg

Rate

0.8

Success Rate

0.6

Rate

0.5

js-pg

Success

0.6

0.4

Success

0.4

chi2-pg

0.3

0.4

0.2

0.2

0.1

0.2

0.0

0

2

4

6

8

0.0

0.5

1.0

1.5

0.0

0.5

1.0

1.5

2.0

2.5

Timesteps

1e6

Timesteps

1e7

Timesteps

1e7

but it can be shown from Lemma 5.1 that fkl -PG converges to the policy that along with maximizing for a 'reward", maximizes the entropy of the state-visitation. A similar result can be shown for χ 2 as well (proof in the Appendix A). This result can be explained by the need for exploration in the larger mazes, hence learning policies to keep the entropy of the state visitation high.

7 Discussion

This paper derives a novel framework for goal-conditioned RL in the form of an on-policy algorithm f -policy gradients which minimizes the f -divergence between the agent's state visitation and the goal distribution. It proves that for certain f -divergences, we can recover the optimal policy while for some, we obtain a policy maximizing the entropy of the state-visitation. Entropy-regularized policy optimization ( s -MaxEnt RL) for metric-based shaping rewards can be shown as a special case of f -PG where f is fkl . f -PG can provide an exploration bonus when the agent has yet not seen the goal. We demonstrated that f -PG can scale up to complex domains.

Through this work, we introduce a new perspective for goal-conditioned RL. By circumventing rewards, f -PG can avoid issues that arise with reward misspecification (Knox et al., 2021). There are several avenues to focus on for future work. First, the current framework is on-policy and poses an exploration challenge. An avenue for future work could be to develop an off-policy way to solve the objective. Second, this paper does not tackle goal distributions with several modes. Such a target distribution would be interesting to tackle in future work.

References

David Abel, Will Dabney, Anna Harutyunyan, Mark K. Ho, Michael L. Littman, Doina Precup, and Satinder Singh. On the expressivity of markov reward. CoRR , abs/2111.00876, 2021. URL https://arxiv.org/abs/2111.00876 .

Jose A. Arjona-Medina, Michael Gillhofer, Michael Widrich, Thomas Unterthiner, Johannes Brandstetter, and Sepp Hochreiter. Rudder: Return decomposition for delayed rewards, 2019.

Andrew G. Barto. Intrinsic Motivation and Reinforcement Learning , pp. 17-47. Springer Berlin Heidelberg, Berlin, Heidelberg, 2013. ISBN 978-3-642-32375-1. doi: 10.1007/978-3-642-32375-1_2. URL https://doi.org/10.1007/978-3-642-32375-1_2 .

Marc G. Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Remi Munos. Unifying count-based exploration and intrinsic motivation, 2016.

Boris Belousov and Jan Peters. f-divergence constrained policy improvement. CoRR , abs/1801.00056, 2018. URL http://arxiv.org/abs/1801.00056 .

Boris Belousov and Jan Peters. Entropic regularization of markov decision processes. CoRR , abs/1907.04214, 2019. URL http://arxiv.org/abs/1907.04214 .

Serena Booth, Julie Shah, Scott Niekum, Peter Stone, and Alessandro Allievi. The perils of trial-anderror reward design: misdesign through overfitting and invalid task specifications. 2023.

Jack Clark and Dario Amodei. Faulty reward functions in the wild, 2016. URL https://openai. com/research/faulty-reward-functions .

Ishan Durugkar, Mauricio Tec, Scott Niekum, and Peter Stone. Adversarial intrinsic motivation for reinforcement learning. CoRR , abs/2105.13345, 2021. URL https://arxiv.org/abs/2105. 13345 .

Ishan Durugkar et al. Estimation and control of visitation distributions for reinforcement learning . PhD thesis, 2023.

Benjamin Eysenbach, Ruslan Salakhutdinov, and Sergey Levine. C-learning: Learning to achieve goals via recursive classification. CoRR , abs/2011.08909, 2020. URL https://arxiv.org/ abs/2011.08909 .

Benjamin Eysenbach, Sergey Levine, and Ruslan Salakhutdinov. Replacing rewards with examples: Example-based policy search via recursive classification. CoRR , abs/2103.12656, 2021. URL https://arxiv.org/abs/2103.12656 .

Alhussein Fawzi, Matej Balog, Aja Huang, Thomas Hubert, Bernardino Romera-Paredes, Mohammadamin Barekatain, Alexander Novikov, Francisco J. R. Ruiz, Julian Schrittwieser, Grzegorz Swirszcz, David Silver, Demis Hassabis, and Pushmeet Kohli. Discovering faster matrix multiplication algorithms with reinforcement learning. Nature , 610(7930):47-53, 2022. doi: 10.1038/s41586-022-05172-4.

Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse reinforcement learning. CoRR , abs/1710.11248, 2017. URL http://arxiv.org/abs/1710.11248 .

Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: datasets for deep data-driven reinforcement learning. CoRR , abs/2004.07219, 2020. URL https://arxiv.org/ abs/2004.07219 .

Theophile Gervet, Soumith Chintala, Dhruv Batra, Jitendra Malik, and Devendra Singh Chaplot. Navigating to objects in the real world. Science Robotics , 8(79):eadf6991, 2023. doi: 10.1126/ scirobotics.adf6991. URL https://www.science.org/doi/abs/10.1126/scirobotics. adf6991 .

Seyed Kamyar Seyed Ghasemipour, Richard S. Zemel, and Shixiang Gu. A divergence minimization perspective on imitation learning methods. CoRR , abs/1911.02256, 2019. URL http://arxiv. org/abs/1911.02256 .

Prasoon Goyal, Scott Niekum, and Raymond J. Mooney. Using natural language for reward shaping in reinforcement learning. CoRR , abs/1903.02020, 2019. URL http://arxiv.org/abs/1903. 02020 .

Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning. CoRR , abs/1910.11956, 2019. URL http://arxiv.org/abs/1910.11956 .

Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. CoRR , abs/1702.08165, 2017. URL http://arxiv.org/abs/1702. 08165 .

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. CoRR , abs/1801.01290, 2018. URL http://arxiv.org/abs/1801.01290 .

Elad Hazan, Sham M. Kakade, Karan Singh, and Abby Van Soest. Provably efficient maximum entropy exploration. CoRR , abs/1812.02690, 2018. URL http://arxiv.org/abs/1812.02690 .

Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. CoRR , abs/1606.03476, 2016. URL http://arxiv.org/abs/1606.03476 .

Rodrigo Toro Icarte, Toryn Klassen, Richard Valenzano, and Sheila McIlraith. Using reward machines for high-level task specification and decomposition in reinforcement learning. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning , volume 80 of Proceedings of Machine Learning Research , pp. 2107-2116. PMLR, 10-15 Jul 2018. URL https://proceedings.mlr.press/v80/icarte18a.html .

Rodrigo Toro Icarte, Ethan Waldie, Toryn Q. Klassen, Richard Anthony Valenzano, Margarita P. Castro, and Sheila A. McIlraith. Learning reward machines: A study in partially observable reinforcement learning. CoRR , abs/2112.09477, 2021. URL https://arxiv.org/abs/2112. 09477 .

John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon Kohl, Andrew Ballard, Andrew Cowie, Bernardino Romera-Paredes, Stanislav Nikolov, Rishub Jain, Jonas Adler, and Demis Hassabis. Highly accurate protein structure prediction with alphafold. Nature , 596:1-11, 08 2021. doi: 10.1038/s41586-021-03819-2.

Leslie Pack Kaelbling. Learning to achieve goals. In IJCAI , pp. 1094-1099. Citeseer, 1993.

Hilbert J. Kappen, Vicenç Gómez, and Manfred Opper. Optimal control as a graphical model inference problem. Machine Learning , 87(2):159-182, feb 2012. doi: 10.1007/s10994-012-5278-7. URL https://doi.org/10.1007%2Fs10994-012-5278-7 .

Liyiming Ke, Matt Barnes, Wen Sun, Gilwoo Lee, Sanjiban Choudhury, and Siddhartha S. Srinivasa. Imitation learning as f-divergence minimization. CoRR , abs/1905.12888, 2019. URL http: //arxiv.org/abs/1905.12888 .

Heecheol Kim, Yoshiyuki Ohmura, and Yasuo Kuniyoshi. Robot peels banana with goal-conditioned dual-action deep imitation learning, 2022.

Bradley Knox, Alessandro Allievi, Holger Banzhaf, Felix Schmitt, and Peter Stone. Reward (mis)design for autonomous driving. CoRR , abs/2104.13906, 2021. URL https://arxiv.org/ abs/2104.13906 .

Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. CoRR , abs/1805.00909, 2018. URL http://arxiv.org/abs/1805.00909 .

Xiao-Yang Liu, Hongyang Yang, Jiechao Gao, and Christina Dan Wang. FinRL. In Proceedings of the Second ACM International Conference on AI in Finance . ACM, nov 2021. doi: 10.1145/ 3490354.3494366. URL https://doi.org/10.1145%2F3490354.3494366 .

Yecheng Jason Ma, Jason Yan, Dinesh Jayaraman, and Osbert Bastani. How far i'll go: Offline goal-conditioned reinforcement learning via f -advantage regression, 2022. URL https://arxiv. org/abs/2206.03023 .

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin A. Riedmiller. Playing atari with deep reinforcement learning. CoRR , abs/1312.5602, 2013. URL http://arxiv.org/abs/1312.5602 .

Ofir Nachum, Yinlam Chow, Bo Dai, and Lihong Li. Dualdice: Behavior-agnostic estimation of discounted stationary distribution corrections. ArXiv , abs/1906.04733, 2019.

Ng, Daishi Harada, and Stuart J. Russell. Policy invariance under reward transformations: Theory and application to reward shaping. In International Conference on Machine Learning , 1999.

Tianwei Ni, Harshit S. Sikchi, Yufei Wang, Tejus Gupta, Lisa Lee, and Benjamin Eysenbach. f-irl: Inverse reinforcement learning via state marginal matching. CoRR , abs/2011.04709, 2020. URL https://arxiv.org/abs/2011.04709 .

Scott Niekum. Evolved intrinsic reward functions for reinforcement learning. Proceedings of the AAAI Conference on Artificial Intelligence , 24(1):1955-1956, Jul. 2010. doi: 10.1609/aaai.v24i1.7772. URL https://ojs.aaai.org/index.php/AAAI/article/view/7772 .

OpenAI, Matthias Plappert, Raul Sampedro, Tao Xu, Ilge Akkaya, Vineet Kosaraju, Peter Welinder, Ruben D'Sa, Arthur Petron, Henrique Pondé de Oliveira Pinto, Alex Paino, Hyeonwoo Noh, Lilian Weng, Qiming Yuan, Casey Chu, and Wojciech Zaremba. Asymmetric self-play for automatic goal discovery in robotic manipulation. CoRR , abs/2101.04882, 2021. URL https: //arxiv.org/abs/2101.04882 .

Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback, 2022.

Matthias Plappert, Marcin Andrychowicz, Alex Ray, Bob McGrew, Bowen Baker, Glenn Powell, Jonas Schneider, Josh Tobin, Maciek Chociej, Peter Welinder, Vikash Kumar, and Wojciech Zaremba. Multi-goal reinforcement learning: Challenging robotics environments and request for research, 2018.

Yury Polyanskiy and Yihong Wu. "Information Theory From Coding to Learning" . Cambridge University Press, 2022. URL https://people.lids.mit.edu/yp/homepage/data/ itbook-export.pdf .

Doina Precup. Eligibility traces for off-policy policy evaluation. Computer Science Department Faculty Publication Series , pp. 80, 2000.

Martin L Puterman. Markov decision processes. Handbooks in operations research and management science , 2:331-434, 1990.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. CoRR , abs/1707.06347, 2017. URL http://arxiv.org/abs/1707. 06347 .

Dhruv Shah, Benjamin Eysenbach, Gregory Kahn, Nicholas Rhinehart, and Sergey Levine. Ving: Learning open-world navigation with visual goals. CoRR , abs/2012.09812, 2020. URL https: //arxiv.org/abs/2012.09812 .

David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, Timothy P. Lillicrap, Karen Simonyan, and Demis Hassabis. Mastering chess and shogi by self-play with a general reinforcement learning algorithm. CoRR , abs/1712.01815, 2017. URL http://arxiv.org/abs/ 1712.01815 .

David Silver, Satinder Singh, Doina Precup, and Richard S. Sutton. Reward is enough. Artificial Intelligence , 299:103535, 2021. ISSN 0004-3702. doi: https://doi.org/10.1016/j.artint.2021.103535. URL https://www.sciencedirect.com/science/article/pii/S0004370221000862 .

Satinder Singh, Richard L. Lewis, Andrew G. Barto, and Jonathan Sorg. Intrinsically motivated reinforcement learning: An evolutionary perspective. IEEE Transactions on Autonomous Mental Development , 2(2):70-82, 2010. doi: 10.1109/TAMD.2010.2051031.

Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems , pp. 5026-5033. IEEE, 2012. doi: 10.1109/IROS.2012.6386109.

Ahmed Touati, Amy Zhang, Joelle Pineau, and Pascal Vincent. Stable policy optimization via off-policy divergence regularization. CoRR , abs/2003.04108, 2020. URL https://arxiv.org/ abs/2003.04108 .

Ronald J Williams and Jing Peng. Function optimization using connectionist reinforcement learning algorithms. Connection Science , 3(3):241-268, 1991.

Peter R Wurman, Samuel Barrett, Kenta Kawamoto, James MacGlashan, Kaushik Subramanian, Thomas J Walsh, Roberto Capobianco, Alisa Devlic, Franziska Eckert, Florian Fuchs, et al. Outracing champion gran turismo drivers with deep reinforcement learning. Nature , 602(7896): 223-228, 2022.

Zeyu Zheng, Junhyuk Oh, and Satinder Singh. On learning intrinsic rewards for policy gradient methods. CoRR , abs/1804.06459, 2018. URL http://arxiv.org/abs/1804.06459 .

Brian D. Ziebart. Modeling purposeful adaptive behavior with the principle of maximum causal entropy. 2010.

Brian D. Ziebart, Andrew L. Maas, J. Andrew Bagnell, and Anind K. Dey. Maximum entropy inverse reinforcement learning. In Dieter Fox and Carla P. Gomes (eds.), AAAI , pp. 1433-1438. AAAI Press, 2008. ISBN 978-1-57735-368-3. URL http://dblp.uni-trier.de/db/conf/aaai/ aaai2008.html#ZiebartMBD08 .

Appendix

A Analysis of J ( θ )

In this section, we will present the proofs for all the Lemmas and Theorems stated in Section 5.1.

A.1 Proof for Theorem 5.1

To prove Theorem 5.1, we need the following Lemmas. Lemma A.1 states that among all policies, the optimal policy has the highest state visitation at the goal.

Lemma A.1. Let D be the set of all possible state visitations for the agent following some policy π ∈ Π . Let π ∗ be the optimal goal-conditioned policy. This optimal policy's state-visitation distribution will have the most measure at the goal for all p π ∈ D i.e., π ∗ = ⇒ p π ∗ ( g ) ≥ p π ( g ) , ∀ p π ∈ D .

Proof. Let π ∗ be the optimal policy and p π ∗ be the corresponding state visitation distribution. The reward for the sparse setting is designed as,

Hence the expected return for a policy π is R π is

The return for the optimal policy is maximum among all policies so R π ∗ ≥ R π , ∀ π ∈ Π . This implies p π ∗ ( g ) ≥ p π ( g ) , ∀ p π ∈ D .

Lemma A.2 states that the f-divergence between p π ( s ) and p g ( s ) is a decreasing function with respect to p π ( s ) . This means that as the objective J ( θ ) obtains its minimum value when p π ( g ) is highest.

Lemma A.2. D f ( p π ( · ) || p g ( · )) is a decreasing function with respect p π ( g ) ∀ f if f ′ ( ∞ ) is defined.

Proof. The goal distribution is assumed to be a Dirac distribution i.e., p g ( s ) = 1 if s = g and 0 everywhere else. The f -divergence between the agent state-visitation distribution, p π and the goal distribution, p g can be defined as,

Let F = D f ( p π || p g ) . Differentiating F w.r.t. p π ( g ) , we get F ′ = f ′ ( p π ( g )) -f ′ ( ∞ ) . Since f is a convex function (by the definition of f -divergence), f ′ ( x ) ≤ f ′ ( y ) , ∀ x ≤ y .

Hence, if f ′ ( ∞ ) is defined, F ′ ≤ 0 . Hence F = D f ( p π || p g ) is a decreasing function with respect p π ( g ) .

Additionally, we need Lemma A.3 and Corollary 1 to complete the proof of Theorem 5.1.

Lemma A.3. If any two policies π 1 and π 2 have the same state visitation at a given goal, they have the same returns for that goal.

Proof. Follows directly from the definition of returns. R π = ❊ p π [ r ( s )] = p π ( g ) . Hence two policies π 1 and π 2 with the same state visitation at the goal will have the same returns.

Corollary 1. Any policy that can lead to the state-visitation distribution of the optimal policy p π ∗ is optimal.

Proof. Directly follows from Lemma A.3.

Theorem 5.1. The policy that minimizes D f ( p π || p g ) for a convex function f with f (1) = 0 and f ′ ( ∞ ) being defined, is the optimal policy.

Proof. Lemma A.1 proves that the optimal policy has the maximum state-visitation probability. Lemma A.2 proves that the f -divergence objective decreases with increasing the state-visitation probability at the goal. In other words, to minimize the f -divergence, we need to maximize the state visitation at goal. Corollary 1 further indicates that any policy that can lead to the state-visitation distribution of the optimal policy i.e., any policy that maximizes the state-visitation distribution at the goal state is an optimal policy.

A.2 Proof for Lemma 5.1

Lemma 5.1. fkl -PG produces a policy that maximizes the reward log p g ( s ) along with the entropy of the state-visitation distribution.

Proof. For fkl -PG, f = u log u . Hence, J ( θ ) = D f ( p π || p g ) can be written as,

❊ where H ( p π ) is the entropy of the agent's state visitation distribution. Minimizing D f ( p π || p g ) will correspond to maximizing the reward r ( s ) = log p g ( s ) and the entropy of p π .

A similar result could be proved for χ 2 divergence:

LemmaA.4. If f ( u ) = ( u -1) 2 ( χ 2 divergence), D f ( p π || p g ) is the upper bound of D FKL ( p π || p g ) -1 . Hence minimizing D χ 2 will also minimize D FKL recovering the entropy regularized policy.

Proof. With f = ( u -1) 2 , D f ( p π || p g ) can be written as,

Since, x > log x ,

Minimizing LHS will also minimize RHS. RHS is essentially D KL ( p π || p g ) -1 . The -1 will not have any effect on the minimization of D KL ( p π || p g ) .

B Gradient based optimization

B.1 Derivation of gradients

Theorem 4.1. The gradient of J ( θ ) as defined in Equation 4 is given by,

Proof. We follow the proof from Ni et al. (2020). Lets start with the state-visitation distribution. In Section 2, it was shown that the state-visitation distribution can be written as,

where f ( s ) = ∫ p ( τ ) e ∑ T t =1 log π θ ( s t ) η τ ( s ) dτ and Z = T ∫ p ( τ ) e ∑ T t =1 log π θ ( s t ) dτ . Differentiating w.r.t. π θ ( s ∗ ) ,

Computing dp θ ( s ) π θ ( s ∗ ) using df ( s ) dπ θ ( s ∗ ) and dZ dπ θ ( s ∗ ) ,

Now we can compute dp θ ( s ) dθ ,

and,

The objective L ( θ ) = D f ( p θ ( s ) || p g ( s )) = ∫ p g ( s ) f ( p θ ( s ) p g ( s ) ) ds The gradient for L ( θ ) will be given by,

Theorem 4.2. Updating the policy using the gradient maximizes ❊ p θ [ η τ ( g )] .

Proof. In goal-based setting, p g is sparse, so we need to use the full definition of f -divergence, D f ( p θ || p g ) = ∑ p g > 0 [ p g ( s ) f ( p θ ( s ) p g ( s ) ) ] + f ′ ( ∞ ) p θ [ p g = 0] = ∑ p g > 0 [ p g ( s ) f ( p θ ( s ) p g ( s ) ) ] + f ′ ( ∞ )(1 -p θ ( g )) . Differentiating with respect to θ gives,

The gradient has two terms, the first term ( f ′ ( p θ ( g )) -f ′ ( ∞ ) ) weighs the gradient based on the value of p θ ( g ) and is always negative. It acts as an adaptive learning schedule, reducing its magnitude as the p θ ( g ) increases. The second term is the gradient of ❊ p θ [ η τ ( g )] . Hence using ∇ θ L ( θ ) , we minimize L ( θ ) which would imply maximizing ❊ p θ [ η τ ( g )] .

B.2 Practical Algorithm

As mentioned in Section 4, the derived gradient is highly sample inefficient. We employ established methods to improve the performance of policy gradients like importance sampling.

The first modification is to use importance sampling weights to allow sampling from previous policy θ ′ . The gradient now looks like,

To reduce the variance in the gradients, the objective can be modified to use the causal connections in the MDP and ensure that the action taken at step t only affects rewards at times t ′ → [ t, T ] . Moreover, a discount factor γ is used to prevent the sum ∑ T t ′ = t f ′ ( p θ ( s t ) p g ( s t ) ) from exploding.

Additionally, the expectation is modified to be over states rather than trajectories,

This gradient computation is still inefficient, because even though the samples are from a previous policy π θ ′ , it still needs to compute ∑ T t ′ = t γ t ′ f ′ ( p θ ( s t ) p g ( s t ) ) , requiring iteration through full trajectories.

We can add a bias to the gradient by modifying f ′ ( p θ ( s t ) p g ( s t ) ) to f ′ ( p θ ′ ( s t ) p g ( s t ) ) in the objective. To ensure the bias is small, an additional constraint needs to be added to keep θ ′ close to θ . Following the literature from natural gradients, the constraint we add is D KL ( p θ ′ || p θ ) . Proximal Policy Optimization (Schulman et al., 2017) showed that in practical scenarios, clipped objective can be enough to do away with the KL regularization term. The final objective that we use is,

B.3 Discounted State-Visitations

The state-visitation distribution defined so far has not considered a discount factor. To include discounting, the state-visitation frequency gets modified to η τ ( s ) = ∑ T t =1 γ t ✶ s t = s . Throughout the derivation of the gradient, we used ∫ η τ ( s ) f ( s ) ds = ∑ T t =1 f ( s t ) but this will be modified to ∫ η τ ( s ) f ( s ) ds = ∑ T t =1 γ t f ( s t ) . the corresponding gradient will be,

This gradient can be modified as before,

Adding importance sampling to the gradient in Equation 13 will give a gradient very similar to Equation 10. In fact, Equation 10 is a biased estimate for the gradient of the f -divergence between the discounted state-visitation distribution and the goal distribution. We can use either of the two gradients but Equation 10 will be preferred for long horizon tasks.

C Gridworld Experiments

C.1 Description of the task

The task involves navigating a gridworld to reach the goal state which is enclosed in a room. The agent can move in any of the four directions and has no idea where the goal is. It needs to explore the gridworld to find the path around the room to reach the goal. The task is further elaborated in Figure 7. The green square represents the agent position while the red square represents the goal.

State: The state visible to the policy is simply the normalized x and y coordinates.

Action: The action is discrete categorical distribution with four categories one for each - left, top, right and bottom.

Reward: The task reward is 1 at the goal and 0 everywhere else. f -PG does not require rewards but the baselines use task rewards.

C.2 Performance of f -PG

In Section 5.1, we had compared fkl -PG and rkl -PG with AIM (Durugkar et al., 2021) and GAIL (Ho & Ermon, 2016). In Figure 8 we present additional baselines AIRL (Fu et al., 2017) and f-AIRL (Ghasemipour et al., 2019).

D Visualizing the learning signals

D.1 Description of the task

To visualize the learning signals, we use the Reacher environment (Figure 9) (Todorov et al., 2012). The task involves rotating a reacher arm (two joints with one end fixed). The applied actions (torques) would rotate the arm so that the free end reaches the goal. The goal is fixed to be at ( -0 . 21 , 0) and the goal distribution is a normal centred at the goal with a standard deviation of 0 . 02 .

(a) FKL

(b) RKL

(c) AIM

(e) AIRL

(f) FAIRL

9

target

target

target

agent

agent

agent

8

8

8

/

6

6

5

3

2

O

8

2

8

2

3

8

target

9

target

target

agent

agent

agent

8

8

8

7

6

6

6

5

5

4

3

3

2

1

State: The state of the original environment contains several things but here we simplify the state space to simply be the position of the free end and the target or the goal position.

Actions: The actions are two dimensional real numbers in [ -1 , 1] which correspond to the torques applied on the two joint respectively.

Reward: The reward is sparse i.e., 1 when the goal is reached by the tip of the arm. But f -PG does not use rewards for training policies.

D.2 Comparing different f -PG

Figure 10 shows the evolution of -f ′ ( p θ ( s ) p g ( s ) ) for the environment. The red regions correspond to signals having a lower value while the darker blue regions correspond to signals with high value. For fkl , the scale of these signals generally varies from 10 to -5 while for χ 2 , the scale varies from 600 to 50 . Also, for the same objective, as the policy trains, these scales generally get smaller in magnitude.

The following can be observed from these plots:

In all the cases, the signals are maximum at the the goal pulling the state-visitations towards the goal.

All of these also push for exploration. This is most pronounced in fkl and χ 2 . These provide significant push towards the unexplored regions which show their inclination towards entropy-maximization, confirming the theory (Lemma 5.1).

E PointMaze experiments

PointMaze (Fu et al., 2020) are continuous state-space domains where the agent needs to navigate to the goal in the 2D maze. The agent and the goal are spawned at a random location in the maze for every episode. There are three levels based on the difficulty of the maze as shown in Figure 11.

State: The state consists of the agent's 2D position and the velocity in the maze. The goal position is appended to the state.

Action: The actions are 2D real numbers in the range [ -1 , 1] correspond to the force applied to the agent in each of the two directions.

Reward: Although f -PG does not use rewards, the baselines use the task reward which is sparse ( 1 when the goal is reached and 0 everywhere else).

田

For the experiments in Section 6.2, the initial and goal states are sampled uniformly over all the 'VALID' states i.e., states that can be reached by the agent. Such an initialization allows discriminatorbased methods to fulfill their coverage assumptions. In Section 6.3, the initialization procedure is modified so that the initial state and the goal state are considerably far. This is done by restricting the sampling of the initial and goal states from disjoint (considerably far away) distributions as shown in Figure 5.

(a) FKL

(b) RKL

(c) JS

(d)

χ

2

D

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies goal-conditioned reinforcement learning and proposes a framework based on f-divergence. The authors show that minimizing the f-divergence between the agent’s visitation and the proposed goal distribution leads to the optimal policy and discuss the practical implementations, connections with previous GCRL methods.\n\nSTRENGTHS:\n1. The paper is well written and easy to follow.\n2. The motivation to study the divergence between the state visitation and goal distribution is straightforward, which is good.\n3. The analysis is thorough. The discussion between the proposed framework and previous methods is also interesting.\n\nWEAKNESSES:\n1. The authors state that they are focusing on sparse-reward (1 if the agent reaches the goal and 0 otherwise). However, $p_g(s)$ will become a Dirac distribution and the f-policy gradient objective makes no sense in this case. \n2. As the authors admit, the algorithm and analysis are meaningless in a full Dirac distribution setting. The authors choose to \"add an ϵ probability to all other states\". Will this uniform ϵ addition + f-PG performs better than uniform reward shaping + standard PG with GC-reward? Intuitively, such uniform addition will still lead to sparse training signal.\n3. It is not very clear about the advantage of the proposed framework. In the strict sparse-reward setting, the framework degrades. Even in dense-reward settings, the analysis only states that optimizing the f-PG objective moves to the optimal policy, which is also true for standard PG methods in GCRL that are based on goal-conditioned (dense) rewards only. Are there any efficiency advantage of f-PG?\n\nQUESTIONS:\nSee the weakness section above.\n\nLIMITATIONS:\nThe authors only discussed some future avenues of their work. I suggest the authors to discuss more about the benefit over the shaping rewards GCRL algorithms, and discuss the previous theoretical results established for shaping rewards methods.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-20"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper focuses on goal-conditioned reinforcement learning (RL), where the goal-conditioned reward function is often sparse. Consequently, standard goal-conditioned RL algorithms suffer from the exploration problem, thus requiring the need to shape rewards. The paper proposes a a new formulation that minimizes the $f$-divergence between the policy's state-visitation distribution and the goal distribution. The paper theoretically shows that minimizing such objective yields the optimal policy and provided empirical evaluations on two maze tasks.\n\nSTRENGTHS:\n- As far as I know, the $f$-divergence between state-visitation distribution and goal distribution is a novel contribution.\n- The paper theoretically shows that optimizing their objective yields the optimal policy.\n- Good visualizations on demonstrating the policies' behaviours.\n\nWEAKNESSES:\n- I think $\\gamma$ should not be introduced as it does not really contribute to the idea. Also, since we wish to maximize $\\mathbb{E}_\\theta\\left[\\eta\\right]$, the aim to achieve the task as quickly as possible is there.\n- How is $\\int \\eta_\\tau(s) ds = T$ in theorem 4.3 when we are not assuming fixed horizon $T$---this is based on page 3, on line 131, where the paper assumes that the trajectory ends when the agent reaches the goal or if maximum horizon $T$ is reached.\n\t- It seems like if the paper is using the fixed-horizon setting, then theorem 4.3 will make sense, in terms of being similar to maximizing the expected return.\n\t- Furthermore, I believe we can actually relate the policy error even if the trained policy is sub-optimal. It will be nice if there is some results regarding suboptimality.\n- On page 5, second last paragraph, why do we not just assign equiprobability to the correct goal states, but $\\varepsilon$ to all other states? It seems like we already have such access anyway. Otherwise, is this not a \"coverage\" assumption?\n- If I am not mistaken, under Gaussian policies, $f$-PG with FKL essentially recovers some sort of L2 distance. I think for reaching tasks this is naturally a good and dense reward, which may explain the better performance than other approaches.\n\t- It will be interesting if the paper provides plot describing the correlation of the \"constructed\" reward and the \"natural\" dense reward (i.e. L2 distance between state and goal).\n\t- It will also be nice to show a task aside from maze and analyze whether FKL still performs better in general.\n- The paper should further describe how the state-visitation distribution $p_\\pi$ and the goal distribution $p_g$ are trained, maybe in the appendix.\n\nQUESTIONS:\n- Based on Figure 4, the paper suggests that the maximizing entropy of state-visitation matters.\n\t- When computing the success rate, are we sampling from the policy, or taking the mean of the policy?\n\t- How would PPO with entropy-regularized objective (or SAC) perform?\n- The paper proposes to learn the distributions separately using kernel density estimators (KDE).\n\t- It seems like this is not really the consequence of the theory, but just a choice, is that right?\n\t- If it is indeed a choice, I am not completely sure if the instability of the discriminator plays a role.\n\t- In the related works section, there were few sentences indicating instability due to non-stationary rewards (in particular, under the subsection shaping rewards). Isn't this approach also non-stationary? Shouldn't there be similar instability as other approaches then?\n\n**Possible typos**\n- On page 2, line 52: \"insufficiency\" instead of \"insufficient\"\n- On page 2, line 90: \"learning\" instead of \"Learning\"\n- On page 3, line 94: \"expert's\" instead of \"expert\"\n- On page 3, line 102: \"learning\" instead of \"Learning\"\n- On page 3, line 133: The expectation should also be over the goal $g$\n- On page 3, line 137: The distribution should include $\\mu_0(s_0)$\n- On page 3, equation (1): $p_\\theta(\\tau ; g)$ instead of $p_\\theta(\\tau)$\n- On page 3, equation (2): $(s) d\\tau$ should be $d\\tau$ on the numerator\n- On page 4, line 160: $\\pi$ instead of $p_\\pi$.\n- Figure 2, legend: $fkl-rew$ should be $fkl-new$\n- Appendix, on page 6, line 86: Should be \"Theorem 4.3\" instead of \"Theorem 4.2\"\n\nLIMITATIONS:\n- The paper requires an assumption of knowing the goal function. That is, reward $r(s, a, s';g) = \\mathbb{I}(s' = g)$. Designing such a goal function can be challenging due to hardware constraints in real-life applications, notably in robotics tasks.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes to minimise the f-divergence $J(\\theta) = D_f(p_\\theta(s) \\| p_g(s))$ between the desired state-distribution $p_g(s)$ and the current one $p_\\theta(s)$. Clipped likelihood-ratio policy gradient is used to estimate the gradient. Kernel density estimation used to estimate the state visitations. Experiments on point reaching in grid world and in continuous-state maze are provided.\n\nSTRENGTHS:\nThe paper is straightforward to follow. The idea and motivation are clear.\nOriginality: this niche has been explored in a number of works. The paper proposes a slightly novel approach, but its merit is not clear as it is only tested on simple point grid world and maze environments. In contrast, the baselines have been shown to work on robotics environments and scale quite well.\nQuality: the paper is OK as a first step in validating the idea, but to really make claims about this algorithm, comparisons on harder environments are necessary. Important details are missing, e.g., it is not really clear from the paper how exactly the gradient is estimated and how KDE was used to estimate visitation frequencies, etc.\nClarity: the paper is mostly clear but it can be improved. E.g., the lemmas and theorems in Sec. 4 are obvious. It would be better to show Algorithm box there and provide more experiments.\nSignificance: currently quite low. Basically, experiments use f-KL, which is the well-known MaxEnt RL. So, the benefit of introducing the whole f-divergence story is unclear.\n\nWEAKNESSES:\nIt is unclear how the authors deal with having $p_g(s)$ in the denominator of the ratio, since by definition it is a delta function. Some epsilon-modification is proposed for grid world, but it doesn't seem like it would scale to any reasonably high-dimensional task.\n\nQUESTIONS:\nSee comments above\n\nLIMITATIONS:\nThe paper does not provide a limitations section. Only future work is mentioned in the conclusion, but this is different from limitations. The authors are encouraged to add the limitations section.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a novel framework for goal-conditioned RL called f-Policy Gradients (f-PG), which minimizes the f-divergence between the agent's state visitation and the goal distribution to provide dense learning signals for exploration in sparse reward settings. The paper derives gradients for various f-divergences to optimize this objective and shows that entropy maximizing policy optimization for commonly used metric-based shaping rewards can be reduced to special cases of f-divergences. The paper compares f-PG with standard policy gradient methods on grid world environment as well as the Point Maze environments. It shows that f-PG performs well in all three environments and better than the baseline shaping rewards. The paper also illustrates how the learning signal evolves with the policy and performs an ablation to compare different f-divergences on their performances on the three Point Maze environments.\n\nSTRENGTHS:\nFirstly, a new framework for goal-conditioned RL which provides dense learning signals for exploration in sparse reward setting is proposed. Secondly, it is proved that minimizing the f -divergence (for some divergences) recovers the optimal policy. Furthermore, the analytical gradient for the objective looks very similar to a policy gradient which allow authors to use established methods from the policy gradient literature to come up with an efficient algorithm for goal-conditioned RL. Finally, empirical evidence in experiment part demonstrate that f-Policy Gradients outperform standard policy gradient methods and baseline shaping rewards in challenging environments. Overall, the paper presents a promising approach for addressing the exploration challenge in goal-conditioned RL.\n\nWEAKNESSES:\nIn this paper, utilizing f-divergence to minimize the mismatch of an agent’s goal-conditioned state visitation to this target distribution is the main contribution. Besides, theoretical as well as empirical analysis has indicated the effectiveness of the methodology. However, some weaknesses still exist and remain to be polish up.\n\n(1) Firstly, f-divergence is not a complete metric in math. Some shortcomings or limitations of it are not been discussed. \n\n(2) Secondly, some notations in the paper are not explained clearly ( e.g. “trajectory dependent state visitation” in section 3, ).\n\n(3) Furthermore, Experiment environments which are adopted in the paper seem to be too simple. Effectiveness of the methodology may not illuminated. Furthermore, some well-known methods in sparse reward setting are not compared.\n\n(4) Novelty of the paper seems to be relatively weak.\n\nWhile the main body of the paper is well-written, there is space for improvement. I defer some of my issues in the appendix to \"Questions\".\n\nQUESTIONS:\nQ1: What will be like if replacing f-divergence with some other mathematical complete metrics (i.e. L2-norm or wasserstein distance)? It is suggested to explain the motivation of utilizing f-divergence more.\n\nQ2: How does f-PG compare to some famous baselines such as curiosity-driven[1] method and DIAYN[2]? It is suggested to add some Comparative experiment。\n\nQ3: Ablation studies about different components of the methodology are suggested to added.\n\nQ4: When the state of environment is in the form of image, is there any difference about the computation of f-divergence?\n\nQ5: Some notations in the paper are not explained clearly ( e.g. “trajectory dependent state visitation” in section 3, ).\n\n[1] Burda, Y., Edwards, H., Pathak, D., Storkey, A., Darrell, T., & Efros, A. A. (2018). Large-scale study of curiosity-driven learning. arXiv preprint arXiv:1808.04355.\n\n[2] Eysenbach, B., Gupta, A., Ibarz, J., & Levine, S. (2018). Diversity is all you need: Learning skills without a reward function. arXiv preprint arXiv:1802.06070.\n\nLIMITATIONS:\nIt would be better if the paper conducts more experiments from the aspects of ablation and comparison with other baselines. Whether effective or not in MARL setting may also be a good extension for the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}