SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Proof. Let  $\pi^{*}$  be the optimal policy and  $p_{\pi^{*}}$  be the corresponding state visitation distribution. The reward for the sparse setting is designed as,
$$
r (s) = \left\{ \begin{array}{l l} 1 & s = g, \\ 0 & \text {o t h e r w i s e}. \end{array} \right.
$$
Hence the expected return for a policy  $\pi$  is  $R_{\pi}$  is
$$
\begin{array}{l} R _ {\pi} = \mathbb {E} _ {p _ {\pi}} [ r (s) ] \\ = p _ {\pi} (g). \\ \end{array}
$$
The return for the optimal policy is maximum among all policies so  $R_{\pi^*} \geq R_{\pi}, \forall \pi \in \Pi$ . This implies  $p_{\pi^*}(g) \geq p_{\pi}(g), \forall p_{\pi} \in \mathcal{D}$ .
Lemma A.2 states that the f-divergence between  $p_{\pi}(s)$  and  $p_g(s)$  is a decreasing function with respect to  $p_{\pi}(s)$ . This means that as the objective  $J(\theta)$  obtains its minimum value when  $p_{\pi}(g)$  is highest.
Lemma A.2.  $D_{f}(p_{\pi}(\cdot)||p_{g}(\cdot))$  is a decreasing function with respect  $p_{\pi}(g)\forall f$  if  $f^{\prime}(\infty)$  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_{\pi}$  and the goal distribution,  $p_{g}$  can be defined as,
$$
\begin{array}{l} D _ {f} \left(p _ {\pi} \mid \mid p _ {g}\right) = \sum_ {p _ {g} > 0} \left[ p _ {g} (s) f \left(\frac {p _ {\pi} (s)}{p _ {g} (s)}\right) \right] + f ^ {\prime} (\infty) p _ {\pi} [ p _ {g} = 0 ] \\ = f \left(p _ {\pi} (g)\right) + f ^ {\prime} (\infty) \left(1 - p _ {\pi} (g)\right). \\ \end{array}
$$
Let  $\mathcal{F} = D_f(p_\pi ||p_g)$ . Differentiating  $\mathcal{F}$  w.r.t.  $p_{\pi}(g)$ , we get  $\mathcal{F}' = f'(p_{\pi}(g)) - f'(\infty)$ . Since  $f$  is a convex function (by the definition of  $f$ -divergence),  $f'(x) \leq f'(y), \forall x \leq y$ .
Hence, if  $f'(\infty)$  is defined,  $\mathcal{F}' \leq 0$ . Hence  $\mathcal{F} = D_f(p_\pi || p_g)$  is a decreasing function with respect to  $p_\pi(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  $\pi_1$  and  $\pi_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_{\pi} = \mathbb{E}_{p_{\pi}}[r(s)] = p_{\pi}(g)$ . Hence two policies  $\pi_1$  and  $\pi_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_{\pi^*}$  is optimal.
Proof. Directly follows from Lemma A.3.  $\square$
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
Proof. For  $fkl$ -PG,  $f = u\log u$ . Hence,  $J(\theta) = D_{f}(p_{\pi}||p_{g})$  can be written as,
$$
\begin{array}{l} D _ {f} (p _ {\pi} | | p _ {g}) = \mathbb {E} _ {p _ {\pi}} \left[ \log \frac {p _ {\pi}}{p _ {g}} \right] \\ = - \left[ \mathbb {E} _ {p _ {\pi}} \left[ \log p _ {g} \right] - \mathbb {E} _ {p _ {\pi}} \left[ \log p _ {\pi} \right] \right] \\ = - \left[ \mathbb {E} _ {p _ {\pi}} [ \log p _ {g} ] + \mathcal {H} (p _ {\pi}) \right] \\ \end{array}
$$
where  $\mathcal{H}(p_{\pi})$  is the entropy of the agent's state visitation distribution. Minimizing  $D_{f}(p_{\pi}||p_{g})$  will correspond to maximizing the reward  $r(s) = \log p_g(s)$  and the entropy of  $p_{\pi}$ .
A similar result could be proved for  $\chi^2$  divergence:
Lemma A.4. If  $f(u) = (u - 1)^2 (\chi^2 \text{ divergence}), D_f(p_\pi || p_g)$  is the upper bound of  $D_{FKL}(p_\pi || p_g) - 1$ . Hence minimizing  $D_{\chi^2}$  will also minimize  $D_{FKL}$  recovering the entropy regularized policy.
Proof. With  $f = (u - 1)^2$ ,  $D_f(p_\pi || p_g)$  can be written as,
$$
\begin{array}{l} D _ {f} \left(p _ {\pi} \mid \mid p _ {g}\right) = \int p _ {g} (s) \left(\frac {p _ {\pi} (s)}{p _ {g} (s)} - 1\right) ^ {2} d s \\ = \int p _ {g} (s) \left(\left(\frac {p _ {\pi} (s)}{p _ {g} (s)}\right) ^ {2} - 2 \frac {p _ {\pi} (s)}{p _ {g} (s)} + 1\right) d s \\ = \int p _ {\pi} (s) \frac {p _ {\pi} (s)}{p _ {g} (s)} - 2 p _ {\pi} (s) + p _ {g} (s) d s \\ = \int p _ {\pi} (s) \frac {p _ {\pi} (s)}{p _ {g} (s)} d s - 1 \\ = \mathbb {E} _ {p _ {\pi} (s)} \left[ \frac {p _ {\pi} (s)}{p _ {g} (s)} \right] - 1 \\ \end{array}
$$
Since,  $x > \log x$
$$
\begin{array}{l} \Longrightarrow \mathbb {E} _ {p _ {\pi} (s)} [ x ] > \mathbb {E} _ {p _ {\pi} (s)} [ \log x ] \\ \Longrightarrow \mathbb {E} _ {p _ {\pi} (s)} \left[ \frac {p _ {\pi} (s)}{p _ {g} (s)} \right] > \mathbb {E} _ {p _ {\pi} (s)} \left[ \log \frac {p _ {\pi} (s)}{p _ {g} (s)} \right] \\ \Longrightarrow \mathbb {E} _ {p _ {\pi} (s)} \left[ \frac {p _ {\pi} (s)}{p _ {g} (s)} \right] - 1 > \mathbb {E} _ {p _ {\pi} (s)} \left[ \log \frac {p _ {\pi} (s)}{p _ {g} (s)} \right] - 1 \\ \end{array}
$$
Minimizing LHS will also minimize RHS. RHS is essentially  $D_{KL}(p_{\pi}||p_g) - 1$ . The  $^{-1}$  will not have any effect on the minimization of  $D_{KL}(p_{\pi}||p_g)$ .
B Gradient based optimization
B.1 Derivation of gradients
$$
\nabla_ {\theta} J (\theta) = \frac {1}{T} \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 {8}
$$
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,
$$
\begin{array}{l} p _ {\theta} (s) \propto \int p (\tau) \Pi_ {t = 1} ^ {T} \pi_ {\theta} (s _ {t}) \eta_ {\tau} (s) d \tau \\ \Rightarrow p _ {\theta} (s) \propto \int p (\tau) e ^ {\sum_ {t = 1} ^ {T} \log \pi_ {\theta} (s _ {t})} \eta_ {\tau} (s) d \tau \\ \Longrightarrow p _ {\theta} (s) = \frac {\int p (\tau) e ^ {\sum_ {t = 1} ^ {T} \log \pi_ {\theta} \left(s _ {t}\right)} \eta_ {\tau} (s) d \tau}{\int \int p (\tau) e ^ {\sum_ {t = 1} ^ {T} \log \pi_ {\theta} \left(s _ {t}\right)} \eta_ {\tau} (s) d \tau d s} \\ \Longrightarrow p _ {\theta} (s) = \frac {\int p (\tau) e ^ {\sum_ {t = 1} ^ {T} \log \pi_ {\theta} \left(s _ {t}\right)} \eta_ {\tau} (s) d \tau}{\int p (\tau) e ^ {\sum_ {t = 1} ^ {T} \log \pi_ {\theta} \left(s _ {t}\right)} \int \eta_ {\tau} (s) d s d \tau} \\ \Longrightarrow p _ {\theta} (s) = \frac {\int p (\tau) e ^ {\sum_ {t = 1} ^ {T} \log \pi_ {\theta} (s _ {t})} \eta_ {\tau} (s) d \tau}{T \int p (\tau) e ^ {\sum_ {t = 1} ^ {T} \log \pi_ {\theta} (s _ {t})} d \tau} \\ \Longrightarrow p _ {\theta} (s) = \frac {f (s)}{Z} \\ \end{array}
$$
where  $f(s) = \int p(\tau)e^{\sum_{t=1}^{T}\log\pi_{\theta}(s_t)}\eta_{\tau}(s)d\tau$  and  $Z = T\int p(\tau)e^{\sum_{t=1}^{T}\log\pi_{\theta}(s_t)}d\tau$ .
Differentiating w.r.t.  $\pi_{\theta}(s^{*})$
$$
\frac {d f (s)}{d \pi_ {\theta} (s ^ {*})} = \frac {\int p (\tau) e ^ {\sum_ {t = 1} ^ {T} \log \pi_ {\theta} (s _ {t})} \eta_ {\tau} (s) \eta_ {\tau} (s ^ {*}) d \tau}{\pi_ {\theta} (s ^ {*})}
$$
$$
\begin{array}{l} \frac {d Z}{d \pi_ {\theta} (s ^ {*})} = \frac {T \int p (\tau) e ^ {\sum_ {t = 1} ^ {T} \log \pi_ {\theta} (s _ {t})} \eta_ {\tau} (s ^ {*}) d \tau}{\pi_ {\theta} (s ^ {*})} \\ = \frac {T f (s ^ {*})}{\pi_ {\theta} (s ^ {*})} \\ \end{array}
$$
Computing  $\frac{dp_{\theta}(s)}{\pi_{\theta}(s^{*})}$  using  $\frac{df(s)}{d\pi_{\theta}(s^{*})}$  and  $\frac{dZ}{d\pi_{\theta}(s^{*})}$ ,
$$
\begin{array}{l} \frac {d p _ {\theta} (s)}{\pi_ {\theta} \left(s ^ {*}\right)} = \frac {Z \frac {d f (s)}{d \pi_ {\theta} \left(s ^ {*}\right)} - f (s) \frac {d Z}{d \pi_ {\theta} \left(s ^ {*}\right)}}{Z ^ {2}} \\ = \frac {\int p (\tau) e ^ {\sum_ {t = 1} ^ {T} \log \pi_ {\theta} (s _ {t})} \eta_ {\tau} (s) \eta_ {\tau} (s ^ {*}) d \tau}{Z \pi_ {\theta} (s ^ {*})} - \frac {f (s)}{Z} T \frac {f (s ^ {*})}{Z \pi_ {\theta} (s ^ {*})} \\ = \frac {\int p (\tau) e ^ {\sum_ {t = 1} ^ {T} \log \pi_ {\theta} (s _ {t})} \eta_ {\tau} (s) \eta_ {\tau} (s ^ {*}) d \tau}{Z \pi_ {\theta} (s ^ {*})} - \frac {T}{\pi_ {\theta} (s ^ {*})} p _ {\theta} (s) p _ {\theta} (s ^ {*}) \\ \end{array}
$$
Now we can compute  $\frac{dp_{\theta}(s)}{d\theta}$ ,
$$
\begin{array}{l} \frac {d p _ {\theta} (s)}{d \theta} = \int \frac {d p _ {\theta} (s)}{\pi_ {\theta} (s ^ {*})} \frac {d \pi_ {\theta} (s ^ {*})}{d \theta} d s ^ {*} \\ = \int \left(\frac {\int p (\tau) e ^ {\sum_ {t = 1} ^ {T} \log \pi_ {\theta} (s _ {t})} \eta_ {\tau} (s) \eta_ {\tau} (s ^ {*}) d \tau}{Z \pi_ {\theta} (s ^ {*})} - \frac {T}{\pi_ {\theta} (s ^ {*})} p _ {\theta} (s) p _ {\theta} (s ^ {*})\right) \frac {d \pi_ {\theta} (s ^ {*})}{d \theta} d s ^ {*} \\ \end{array}
$$
$$
\begin{array}{l} = \frac {1}{Z} \iint p (\tau) e ^ {\sum_ {t = 1} ^ {T} \log \pi_ {\theta} (s _ {t})} \eta_ {\tau} (s) \eta_ {\tau} (s ^ {*}) \frac {\nabla_ {\theta} \pi_ {\theta} (s ^ {*})}{\pi_ {\theta} (s ^ {*})} d s ^ {*} d \tau - T p _ {\theta} (s) \int p _ {\theta} (s ^ {*}) \frac {\nabla_ {\theta} \pi_ {\theta} (s ^ {*})}{\pi_ {\theta} (s ^ {*})} d s ^ {*} \\ = \frac {1}{Z} \int p (\tau) e ^ {\sum_ {t = 1} ^ {T} \log \pi_ {\theta} (s _ {t})} \eta_ {\tau} (s) \sum_ {t = 1} ^ {T} \nabla_ {\theta} \log \pi_ {\theta} (s _ {t}) d \tau - T p _ {\theta} (s) \mathbb {E} _ {s \sim p _ {\theta} (s)} [ \nabla_ {\theta} \log \pi_ {\theta} (s) ] \\ \end{array}
$$
The objective  $L(\theta) = D_{f}(p_{\theta}(s)||p_{g}(s)) = \int p_{g}(s)f\left(\frac{p_{\theta}(s)}{p_{g}(s)}\right)ds$
The gradient for  $L(\theta)$  will be given by,
Theorem 4.2. Updating the policy using the gradient maximizes  $\mathbb{E}_{p_{\theta}}[\eta_{\tau}(g)]$
Proof. In goal-based setting,  $p^g$  is sparse, so we need to use the full definition of  $f$ -divergence,  $D_f(p_\theta ||p_g) = \sum_{p_g > 0} [p_g(s)f(\frac{p_\theta(s)}{p_g(s)})] + f'(\infty)p_\theta[p_g = 0] = \sum_{p_g > 0} [p_g(s)f(\frac{p_\theta(s)}{p_g(s)})] + f'(\infty)(1 - p_\theta(g))$ . Differentiating with respect to  $\theta$  gives,
$$
\begin{array}{l} \nabla_ {\theta} L (\theta) = \left(f ^ {\prime} \left(p _ {\theta} (g)\right) - f ^ {\prime} (\infty)\right) \nabla_ {\theta} p _ {\theta} (s) \\ = \left(f ^ {\prime} \left(p _ {\theta} (g)\right) - f ^ {\prime} (\infty)\right) \left(\frac {1}{T} \int p _ {\theta} (\tau) \eta_ {\tau} (g) \sum_ {t = 1} ^ {T} \nabla_ {\theta} \log \pi_ {\theta} \left(s _ {t}\right) d \tau - T p _ {\theta} (s) \mathbb {E} _ {s \sim p _ {\theta} (s)} [ \nabla_ {\theta} l o g \pi_ {\theta} (s) ]\right) \\ = \left(f ^ {\prime} \left(p _ {\theta} (g)\right) - f ^ {\prime} (\infty)\right) \left(\frac {1}{T} \mathbb {E} _ {\tau \sim p _ {\theta} (\tau)} \left[ \eta_ {\tau} (g) \sum_ {t = 1} ^ {T} \nabla_ {\theta} \log \pi_ {\theta} \left(s _ {t}\right) \right] \right. \\ \left. - p _ {\theta} (g) \mathbb {E} _ {\tau \sim p _ {\theta} (\tau)} \left[ \sum_ {t = 1} ^ {T} \nabla_ {\theta} \log \pi_ {\theta} (s) \right]\right) \\ = \frac {1}{T} \left(f ^ {\prime} \left(p _ {\theta} (g)\right) - f ^ {\prime} (\infty)\right) \mathbb {E} _ {\tau \sim p _ {\theta} (\tau)} \left[ \sum_ {t = 1} ^ {T} \nabla_ {\theta} \log \pi_ {\theta} \left(s _ {t}\right) \eta_ {\tau} (g) \right] \\ \end{array}
$$
The gradient has two terms, the first term  $\left(f'(p_{\theta}(g)) - f'(\infty)\right)$  weighs the gradient based on the value of  $p_{\theta}(g)$  and is always negative. It acts as an adaptive learning schedule, reducing its magnitude as the  $p_{\theta}(g)$  increases. The second term is the gradient of  $\mathbb{E}_{p_{\theta}}[\eta_{\tau}(g)]$ . Hence using  $\nabla_{\theta}L(\theta)$ , we minimize  $L(\theta)$  which would imply maximizing  $\mathbb{E}_{p_{\theta}}[\eta_{\tau}(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  $\theta^{\prime}$ . The gradient now looks like,
$$
\nabla_ {\theta} J (\theta) = \mathbb {E} _ {\tau \sim p _ {\theta^ {\prime}} (\tau)} \left[ \frac {\pi_ {\theta} (\tau)}{\pi_ {\theta^ {\prime}} (\tau)} \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} (s _ {t})}{p _ {g} (s _ {t})}\right) \right] \right]. \tag {9}
$$
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' \to [t, T]$ . Moreover, a discount factor  $\gamma$  is used to prevent the sum  $\sum_{t'=t}^{T} f'\left(\frac{p_{\theta}(s_t)}{p_g(s_t)}\right)$  from exploding.
Additionally, the expectation is modified to be over states rather than trajectories,
$$
\nabla_ {\theta} J (\theta) = \mathbb {E} _ {s _ {t}, a _ {t} \sim p _ {\theta^ {\prime}} \left(s _ {t}, a _ {t}\right)} \left[ \frac {\pi_ {\theta} \left(a _ {t} \mid s _ {t}\right)}{\pi_ {\theta^ {\prime}} \left(a _ {t} \mid s _ {t}\right)} \nabla_ {\theta} \log \pi_ {\theta} \left(a _ {t} \mid s _ {t}\right) \sum_ {t ^ {\prime} = t} ^ {T} \gamma^ {t ^ {\prime}} f ^ {\prime} \left(\frac {p _ {\theta} \left(s _ {t}\right)}{p _ {g} \left(s _ {t}\right)}\right) \right]. \tag {10}
$$
This gradient computation is still inefficient, because even though the samples are from a previous policy  $\pi_{\theta'}$ , it still needs to compute  $\sum_{t'=t}^{T} \gamma^{t'} f'\left(\frac{p_{\theta}(s_t)}{p_g(s_t)}\right)$ , requiring iteration through full trajectories. We can add a bias to the gradient by modifying  $f'\left(\frac{p_{\theta}(s_t)}{p_g(s_t)}\right)$  to  $f'\left(\frac{p_{\theta'}(s_t)}{p_g(s_t)}\right)$  in the objective. To ensure the bias is small, an additional constraint needs to be added to keep  $\theta'$  close to  $\theta$ . Following the literature from natural gradients, the constraint we add is  $D_{KL}(p_{\theta'}||p_{\theta})$ . 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,
$$
\nabla_ {\theta} J (\theta) = \mathbb {E} _ {s _ {t}, a _ {t} \sim p _ {\theta^ {\prime}} \left(s _ {t}, a _ {t}\right)} \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 {11}
$$
where  $r_{\theta}(s_t) = \frac{\nabla_\theta\pi_\theta(a_t|s_t)}{\pi_{\theta'}(s_at|s_t)}$  and  $F_{\theta^{\prime}}(s_{t}) = \sum_{t^{\prime} = t}^{T}\gamma^{t^{\prime}}f^{\prime}\left(\frac{p_{\theta^{\prime}}(s_{t})}{p_{g}(s_{t})}\right)$ .
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  $\eta_{\tau}(s) = \sum_{t=1}^{T} \gamma^t \mathbb{1}_{s_t=s}$ . Throughout the derivation of the gradient, we used  $\int \eta_{\tau}(s) f(s) ds = \sum_{t=1}^{T} f(s_t)$  but this will be modified to  $\int \eta_{\tau}(s) f(s) ds = \sum_{t=1}^{T} \gamma^t f(s_t)$ . The corresponding gradient will be,
$$
\nabla_ {\theta} J (\theta) = \frac {1}{T} \mathbb {E} _ {\tau \sim p _ {\theta} (\tau)} \left[ \left[ \sum_ {t = 1} ^ {T} \gamma^ {t} \nabla_ {\theta} \log \pi_ {\theta} \left(a _ {t} \mid s _ {t}\right) \right] \left[ \sum_ {t = 1} ^ {T} \gamma^ {t} f ^ {\prime} \left(\frac {p _ {\theta} \left(s _ {t}\right)}{p _ {g} \left(s _ {t}\right)}\right) \right] \right]. \tag {12}
$$
This gradient can be modified as before,
$$
\nabla_ {\theta} J (\theta) = \mathbb {E} _ {s _ {t}, a _ {t} \sim p _ {\theta} (s _ {t}, a _ {t})} \left[ \gamma^ {t} \nabla_ {\theta} \log \pi_ {\theta} \left(a _ {t} \mid s _ {t}\right) \sum_ {t ^ {\prime} = t} ^ {T} \gamma^ {t ^ {\prime}} f ^ {\prime} \left(\frac {p _ {\theta} (s _ {t})}{p _ {g} (s _ {t})}\right) \right]. \tag {13}
$$
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.
Figure 7: Description of the gridworld: The bold lines show the walls, green square is the start position and the red square is the goal.
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.
Figure 9: The Reacher environment with fixed goal
(a) FKL
(d) GAIL
Figure 8: Gridworld: The agent needs to move from the green circle to the red circle. The state visitations of the final policies are shown when using our framework for training (fkl, rkl) compared with AIM and GAIL trained on top of soft Q learning.
(b) RKL
(e) AIRL
(c) AIM
(f) FAIRL
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^{\prime}\left(\frac{p_{\theta}(s)}{p_{g}(s)}\right)$  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  $\chi^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:
1. In all the cases, the signals are maximum at the goal pulling the state-visitions towards the goal.
2. All of these also push for exploration. This is most pronounced in  $fkl$  and  $\chi^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).
Figure 11: Description of PointMaze environments: PointMaze-U (left), PointMaze-Medium (middle), PointMaze-Large(right).
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 discriminator-based 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
Figure 10: Evolution of  $-f^{\prime}(\frac{p_{\theta}(s)}{p_g(s)})$  along with the corresponding state-visitations for different  $f$ -divergences -  $fkl$ ,  $rkl$ ,  $js$  and  $\chi^2$ . The scales for the value of these signals are not shown but they vary as the policy converges.  $f$ -PG provides dense signals for pushing towards exploration.
(d)  $\chi^2$