SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The critic is trained with the same loss function  $\mathcal{L}_{SR}$  as defined in Equation 10, where the gradients are not propagated through the target. The actor is optimized to maximize the estimated utility function (Equation 2). Since the actor is deterministic, policy gradients are computed using an adaptation of the deterministic policy gradient theorem [55]. Because the actor network has no dependency on predecessor trace (which depends on the observed states only), gradients for the actor parameters are obtained by applying chain rule leading to the following gradient of Equation 7 (please refer to Proposition 2 for more details on the derivation):
$$
\nabla_ {\mu} J (\pi_ {\mu}) = \mathbb {E} _ {\tau \sim \rho} \left[ \sum_ {i} z _ {i} \nabla_ {\mu} \pi_ {\mu} (\tau) \nabla_ {a} \boldsymbol {\psi} _ {i} (\tau , a) \mid_ {a = \pi_ {\mu} (\tau)} \right], \tag {12}
$$
where  $z_{i} = -\log [\pmb{\eta}(\tau_{: - 1})_{i} + \pmb{\psi}(\tau ,\pi_{\mu}(\tau))_{i}] - 1$  is the multiplicative factor for state  $i$ , and depends on the expected probability of visiting a state. The factor  $z_{i}$  can take values between -1 and  $\infty$ , with positive values of high magnitude for states with low visitation probability and negative values for state with high probability of visitation. Thus, the factor  $z_{i}$  scales the policy gradients to maximize the entropy of the state visitation distribution. In Algorithm 3 we outline the training procedure for the infinite action space framework.
To analyze and test if  $\eta \psi$ -Learning learns an efficient exploration policy, we evaluate the proposed method on a set of discrete and continuous control tasks. In these experiments, we are recording a set of different performance measures to access if the resulting exploration policies do in fact maximize the entropy of the state visitation distribution and if most states are explored by  $\eta \psi$ -Learning. The following results demonstrate that by maintaining a predecessor representation and conditioning the SR on the simulated trajectory prefix, the  $\eta \psi$ -Learning agent learns a deterministic exploration policy that minimizes the number of interactions needed to visit all states. In addition, we expand our method to continuous control tasks and demonstrate how  $\eta \psi$ -Learning can efficiently explore in complex domains with infinite action space. Our method is ideal for searching out rewards in difficult sparse reward environments. We compare  $\eta \psi$ -Learning, which learns to explore an environment as efficiently as possible, to recent meta-RL methods [70] that aim to learn how to optimally explore an environment to infer the rewarding or goal state. Lastly, we show the applicability of  $\eta \psi$ -Learning on standard RL tasks where the extrinsic rewards are sparse and present how the proposed method can be combined with existing algorithms.
**Environments:** We experiment with different tasks with both finite and infinite action spaces. The ChainMDP and RiverSwim [58] is a six-state chain where the transitions are deterministic or stochastic, respectively. In these tasks a Markovian policy cannot cover the state space uniformly because the agent has to pace back and forth along the chain, visiting the same state multiple times. For the RiverSwim environment, a non-stationary policy, a policy that is a function of the time step, cannot optimally cover all states because non-determinism in the transitions can place the agent into different states at random. Furthermore, we include the  $5 \times 5$  grid world example used in Figure 1. We also test  $\eta \psi$ -Learning on two harder exploration tasks—the TwoRooms and FourRooms domains, which are challenging because it is not possible to obtain exact uniform visitation distribution due to the wall structure. For continuous control tasks, we evaluate on Reacher and Pusher tasks, where the
Figure 2: Comparison of  $\eta \psi$ -Learning and MaxEnt [23] on three metrics: Entropy (top row) of state visitation distribution, State Coverage (middle row) representing the fraction of state space visited, and Search Completion Time (bottom row) denoting steps taken to cover the state space.
agent has a robotic-arm with multiple joints. The task is to maximize the entropy over the locations covered by the fingertip of the robotic-arm. Appendix F provides more details on the environments and the hyper-parameters are reported in Appendix G.
Prior Methods: To our knowledge, existing work focuses on learning Markovian exploration policies [43]. We use MaxEnt [23] as a baseline agent for our experiments because this method optimizes a similar maximum entropy objective as  $\eta \psi$ -Learning—with the difference that MaxEnt learns a Markovian policy and resorts to randomization to obtain a close to uniform state visitation distribution. We have also compared with MEPOL [42] on continuous control tasks which leans a Markovian policy and uses kNN-based estimators to compute the entropy. A comparison with SMM [31] is skipped because the method optimizes a similar maximum entropy objective with a Markovian policy and cannot express the same exploration behaviour as  $\eta \psi$ -Learning.
Evaluation Metrics: Entropy measures a method's ability to have similar visitation frequency for each state in the state space. This signifies the gap between the observed state visitation distribution and the optimal distribution that maximizes the entropy term. The Entropy metric is computed using the objective defined in Equation 2 over a single trajectory generated by the agent. A constant discount factor of  $\gamma(t) = \frac{1}{h}$  is used to obtain the state visitation distribution during evaluation. An agent can maximize this measure without actually exploring all states of an environment—a desirable property for RL where rewards may be sparse and hidden in complex to-reach states. We record the state coverage metric which represents the fraction of states in the environment visited by the agent at least once within a trajectory. Lastly, we want agents to explore the state space efficiently. For example, an optimal agent can sweep through the gridworld presented in Figure 1 with a search completion time of 15 steps (Figure 1(b) shows an optimal trajectory). The search completion time metric measures the steps taken to discover each state in the environment. All results report the mean performance computed over 5 random seeds with  $95\%$  confidence intervals shading.
Quantitative Results: Figure 2 presents the results obtained for  $\eta \psi$ -Learning and MaxEnt [23]. Compared to MaxEnt, which learns a Markovian policy,  $\eta \psi$ -Learning achieves  $20 - 50\%$  higher entropy. This indicates that the MaxEnt algorithm by learning a Markovian and stochastic policy was randomizing at certain states which lead to sub-optimal behaviors. The performance gain was more prominent in grid-based environments because the MaxEnt agent was visiting some states more frequently than others, which are harder to explore efficiently. Furthermore, high entropy values suggest that the agent visits states with similar frequency in the environment and does not get stuck at a particular state. We attribute this behavior of  $\eta \psi$ -Learning to the proposed Q-function that picks action to visit different states and maximize the objective.
Figure 3: (a) Heatmap of state visitation distribution by unrolling a trajectory using MaxEnt (left) and  $\eta \psi$ -Learning (right) on TwoRooms and FourRooms environments. (b) Visualization of learned SR of each action (denoted with  $\psi(.)$ ) at time steps  $T = 7, 14$  for a trajectory using  $\eta \psi$ -Learning on  $5 \times 5$  grid.  $(s, a)$  denotes the state (black) and action taken by the agent (direction of white arrow),  $\eta$  is the predecessor representation till time  $T$  (higher values have darker shade)
Figure 2 also shows that  $\eta \psi$ -Learning achieves optimal state coverage across environments exemplifying that  $\eta \psi$ -Learning while maximizing the entropic measure also learns to cover the state space within a single trajectory. However, the baseline MaxEnt was not able to discover all the states in the environment. MaxEnt was unable to visit all the states in ChainMDP and RiverSwim environments with trajectory length of 20 and 50, respectively. Moreover, the state coverage of MaxEnt was around  $50 - 60\%$  on the harder TwoRooms and FourRooms tasks, where the agent has to navigate between different rooms and is required to remember the order of visiting different rooms. These results reveal that Markovian policy limits an agent's ability to maximize the state coverage in a task. The proposed method also outperformed the baseline on the search completion time metric across all environments. Notably, on ChainMDP, the  $5 \times 5$  gridworld, and TwoRooms environments,  $\eta \psi$ -Learning converged within 500 episodes. However,  $\eta \psi$ -Learning did not achieve optimal search exploration time on the FourRooms environment as it missed a spot in a room and resorted to it later in the episode.
To further understand the gains offered by  $\eta \psi$ -Learning, we visualized the state visitation distributions on a single trajectory (in Figure 3(a)). On the TwoRooms environment,  $\eta \psi$ -Learning had similar density on the states in both the rooms, where the density is more around the center. This is because the agent was sweeping across rooms alternatively.  $\eta \psi$ -Learning showed a better-visited state distribution on the FourRooms environment with more distributed density across states. However, MaxEnt was not visiting all states and also visited a few states more frequently than others, elucidating the lower performance on entropy and state coverage. We further visualized the learned SR to see if  $\eta \psi$ -Learning learns a SR for the optimal exploration policy through generalized policy improvement [3]. For this analysis, we sampled a trajectory on  $5 \times 5$  gridworld. Figure 3(b) reports the heatmaps of the learned SR vector for each action at different steps in the trajectory. We observe that the SR vector for each action has lower density on the states already observed in the trace. This exemplifies that the learned SR captures the history of the visited states that further aids in taking actions to maximize the entropy of state visitation distribution. We also study if MaxEnt can show similar gains when trained with a recurrent policy (Appendix I.1) and compared the agents when evaluated across multiple trajectories (Appendix I.2).
Continuous Control tasks: The efficacy of  $\eta \psi$ -Learning is further tested on environments with infinite action space. Figure 4(a) reports the Entropy and State Coverage metric on Reacher and Pusher environments, where  $\eta \psi$ -Learning outperformed the baseline MaxEnt on both metrics. The gains are more significant on the Pusher environment which is a harder task because of multiple hinges in the robotic-arm. The proposed method  $\eta \psi$ -Learning achieves close to  $90\%$  coverage in both environments, whereas the MaxEnt had only close to  $50\%$  and  $40\%$  coverage on Reacher and Pusher environments, respectively. In Figure 4(b), heatmaps of the state visitation density for a single trajectory shows that  $\eta \psi$ -Learning has more uniformly distributed density compared to MaxEnt. The Pusher environment has high density at the top-right corner of the grid denoting the time taken by the agent to move the fingertip to other locations from the starting state. Notably, the proposed method  $\eta \psi$ -Learning has lower density at the starting state and we believe that conditioning on the history of states is guiding the agent to move the robotic-arm to other locations to maximize the entropy over the state space. In Appendix I.5, a visualization of a rolled-out trajectory generated using  $\eta \psi$ -Learning is
(c)
Figure 4: (a) Comparison of  $\eta \psi$ -Learning with Random policy, MaxEnt [43], and MEPOL [42] on Reacher and Pusher tasks (b) Heatmaps of visitation distribution of MaxEnt (top) and  $\eta \psi$ -Learning (bottom), (c) comparison with VariBAD [70] on State Coverage and Goal Search Time metrics.
presented showing that the agent learns to efficiently maneuver the fingertip of the robotic-arm to different locations in the environment.
Comparison with Meta-RL: A question central to Meta-RL [16, 70, 37] is the ability to quickly explore a task and find rewarding states in complex tasks where the rewards are sparse. In this context, Zintgraf et al. [70] present the VariBAD method, which maintains a belief over different tasks to infer the optimal policy—leading to efficient exploration behaviour that enables the agent to discover rewarding states quickly. Similar to VariBAD, the predecessor representation  $\eta$  in  $\eta \psi$ -Learning keeps track of which states have been explored and which states are not explored. Figure 4(c) compares the exploration behaviour of  $\eta \psi$ -Learning to VariBAD: In terms of the State Coverage and Goal Search Time metric,  $\eta \psi$ -Learning outperforms VariBAD significantly because  $\eta \psi$ -Learning is designed to optimize the entropy of the state visitation frequencies of a single trajectory instead of performing Bayes-adaptive inference across a task space. We refer the reader to Appendix H for more details on this experiment.
$\eta \psi$ -Learning in Sparse Reward Tasks: Sparse reward environments pose a challenge where an agent has to discover the reward function by visiting different regions of the state space. Through this experiment, we demonstrate how an agent can leverage  $\eta \psi$ -Learning as an exploratory bonus and improve its efficiency in such tasks. The experiment is conducted on Sparse Mountain-Car environment where the agent receives a positive reward upon reaching the goal position. This is hard because the agent needs to plan to visit different positions in the environment. We have used TD3 [17] as the baseline algorithm for this task. To compare with leading exploration methods, we add TD3 combined with Count-Based [6] and First-Occupancy based [40] to the baselines. It is to be noted that we have used these methods as episodic bonuses as described in [67]. We also propose a variant of TD3 equipped with the proposed method and call it TD3- $\eta \psi$ -Learning, where we learn two critics to estimate the sum of extrinsic rewards and the entropy of state visitation distribution (Equation 7). Appendix J describes the experimental setup in more detail and Algorithm 4 provides the pseudocode of the proposed variant. The evaluation is done across two metrics- Average Return and Average Episode Length denoting the steps taken to reach the goal state. Figure 5 presents the results where the proposed method outperforms the baseline algorithms and is sample-efficient at learning the task.
TD3 combined with First-Occupancy based intrinsic reward performed better than Count-Based bonus but still takes twice number of steps as the proposed method to solve the task. This shows that the proposed method can improve sample efficiency in standard RL tasks, especially with sparse
Figure 5: Results on Sparse Mountain Car environment.
rewards. We leave scaling the proposed method with leading architectures [21, 25, 53] to more challenging tasks for future research.
To explore efficiently, an intelligent agent needs to consider past episodic experiences to decide on the next exploratory action. We demonstrate how the predecessor representation—an encoding of past state visitations—can be combined with the successor representation—a prediction of future state visitations—to learn efficient exploration policies that maximize the state-visitation-distribution entropy. Across a set of different environments, we illustrate how  $\eta \psi$ -Learning consistently reasons across different trajectories to explore near optimally—a task that is NP-hard [43].
To the best of our knowledge,  $\eta \psi$ -Learning is the first algorithm that combines predecessor and successor representations to estimate the state visitation distribution. Furthermore,  $\eta \psi$ -Learning learns a non-Markovian policy and can therefore express exploration behavior not afforded by existing methods [23, 31, 42, 18]. To further increase the applicability of  $\eta \psi$ -Learning, one interesting direction of future research is to extend  $\eta \psi$ -Learning to POMDP environments where states are either partially observable or complex such as images. This is challenging because the agent has to learn an embedding of state observations that capture only the relevant components of the state space to maximize the entropy. We believe a promising approach would be to leverage the idea of Successor Measures (SM) [62, 63, 15] which have shown promising results when scaled to high-dimensional inputs like images. Furthermore, the presented approach can be also used for designing other algorithms that control the state visitation distribution. An application is goal-conditioned RL, where the agents need to minimize the KL divergence between visitation distribution of policy and goal-distribution [31, 47]. Another application is Safe RL [66] where agents receive a penalty upon visiting unsafe states to avoid observing them.
We study reinforcement learning, which aims to enable autonomous agents to acquire search behaviors. This study of developing exploration behaviors in reinforcement learning is guided by a fundamental curiosity about the nature of autonomous learning; it has a number of potential practical applications and broad implications. First, autonomous exploration for pre-training in general, can enable autonomous agents to acquire useful skills with less human intervention and effort, potentially improving the feasibility of learning-enabled robotic systems. Second, the practical applications that we illustrate, such as applications to continuous environments, can accelerate reinforcement learning in certain settings. Specific to our method, finite length entropy maximization may also in the future offer a useful tool for search and rescue, by equipping agents with an objective that causes them to explore a space systematically to locate lost items. However, these types of reinforcement learning methods also have a number of uncertain broad implications: agents that explore the environment and attempt to acquire open-ended skills may carry out unexpected or unwanted behaviors, and would require suitable safety mechanisms of their own during training.
The authors would like to thank Harley Wiltzer for his valuable feedback and discussions. The writing of the paper also benefited from discussions with Darshan Patil, Chen Sun, Mandana Samiei, Vineet Jain, and Arushi Jain. The authors want to acknowledge funding support from NSERC and CIFAR, AJ and IR acknowledge the support from Canada CIFAR AI Chair Program and from the Canada Excellence Research Chairs (CERC) program. The authors are also grateful to Mila (mila.quebec) IDT, NVIDIA and Digital Research Alliance of Canada for computing resources.
[1] Majid Abdolshah, Hung Le, Thommen Karimpanal George, Sunil Gupta, Santu Rana, and Svetha Venkatesh. A new representation of successor features for transfer across dissimilar environments. In International Conference on Machine Learning, pages 1-9. PMLR, 2021.
[2] Duncan Bailey and Marcelo Mattar. Predecessor features. arXiv preprint arXiv:2206.00303, 2022.
A Convergence Analysis
To gain a deeper understanding why the  $\eta \psi$ -Learning converges to a maximum entropy policy, we consider in this section a simplified dynamic programming variant in Algorithm 1. Note that the  $\eta \psi$ -Learning estimates the SR for a finite CMP for a finite horizon length  $h$ . Consequently, the trajectory-action conditioned SR  $\psi^{\pi}(\tau :T,a)$  and exploration policy  $\pi$  can be stored in an exponentially large but finite look-up table. Furthermore, with every transition an additional state is appended to the trajectory  $\tau :T$ , meaning the agent cannot loop back to the same trajectory. Using these two properties, we state a dynamic programming variant of  $\eta \psi$ -Learning in Algorithm 1 and then prove its convergence to a policy that maximizes the entropy term  $H\left(\eta (\tau :T - 1) + \psi^{\pi}(\tau :T,a)\right)$  at every time step.
Algorithm 1  $\eta \psi$ -Learning: Dynamic Programming Framework
1: for all  $\tau_h, a$  do
2:  $\pmb{\psi}^{\pi}(\tau_{h},a)\gets \pmb{e}_{s_h}$
3: end for
4: for  $t = h, \dots, 2$  do
5: for all  $\tau_{:t}, a$  do
6:  $\pi (\tau_{:t})\gets \arg \max_{a}H\left(\pmb {\eta}(\tau_{:t - 1}) + \pmb{\psi}^{\pi}(\tau_{:t},a)\right)$
7:  $\pmb{\psi}^{\pi}(\tau_{:t - 1},a)\gets \pmb{e}_{s_{t - 1}} + \gamma (t)\pmb{\psi}^{\pi}(\tau_{:t},\pi (\tau_{:t}))$
8: end for
9: end for
10: return  $\pi$  such that  $\pi(\tau) = \arg \max_{a} H(\pmb{\eta}(\tau_{-1}) + \pmb{\psi}(\tau, a))$ .
The convergence proof uses the following property of the predecessor trace  $\pmb{\eta}$  and SR  $\psi^{\pi}$ : Consider a trajectory  $\tau$  which selects action  $a_{T}$  at time step  $T$ , then
$$
\begin{array}{l} \boldsymbol {\eta} (\tau ; T - 1) + \boldsymbol {\psi} ^ {\pi} (\tau ; T, a _ {T}) = \underbrace {\sum_ {t = 1} ^ {T - 1} \gamma (t) \boldsymbol {e} _ {s _ {t}}} _ {= \boldsymbol {\eta} (\tau ; T - 1)} + \underbrace {\mathbb {E} _ {\tau_ {T + 1 : , \pi}} \left[ \sum_ {t = T} ^ {h} \gamma (t) \boldsymbol {e} _ {s _ {t}} \mid \tau ; T , a _ {T} \right]} _ {= \boldsymbol {\psi} ^ {\pi} (\tau ; T, a _ {T})} (byEq.5) \\ = \sum_ {t = 1} ^ {T - 1} \gamma (t) \boldsymbol {e} _ {s _ {t}} + \mathbb {E} _ {\tau_ {T + 1}; \pi} \left[ \gamma (T) \boldsymbol {e} _ {s _ {T}} + \sum_ {t = T + 1} ^ {h} \gamma (t) \boldsymbol {e} _ {s _ {t}} \Bigg | \tau_ {: T}, a _ {T} \right] (13) \\ = \sum_ {t = 1} ^ {T} \gamma (t) \boldsymbol {e} _ {s _ {t}} + \mathbb {E} _ {\tau_ {T + 1:}, \pi} \left[ \sum_ {t = T + 1} ^ {h} \gamma (t) \boldsymbol {e} _ {s _ {t}} \mid \tau_ {: T}, a _ {T} \right] (14) \\ = \boldsymbol {\eta} (\tau_ {: T}) + \mathbb {E} _ {\tau_ {T + 1}:, \pi} \left[ \boldsymbol {\psi} ^ {\pi} \left(\tau : T + 1, \pi \left(\tau : T + 1\right)\right) \mid \tau : T, a _ {T} \right] (15) \\ \end{array}
$$
Using this identity, we can prove the convergence of Algorithm 1.
Proposition 1. The policy  $\pi^{*}$  returned by Algorithm 1 is such that for every  $t$ -step trajectory  $\tau_{:t}$  where  $t \leq h$ ,
$$
\pi^ {*} (\tau_ {: t}) \in \arg \max  _ {\pi} H (\boldsymbol {\eta} (\tau_ {: t - 1}) + \boldsymbol {\psi} ^ {\pi} (\tau_ {: t}, \pi (\tau_ {: t}))). \tag {16}
$$
Proof. The proof proceeds by induction on the length of an  $h$ -step trajectory, starting with a length of  $h$  and iterating to a length of one.
Induction hypothesis: We define a sub-sequence optimal policy  $\pi_t$  such that for every  $k$ -step trajectory prefix  $\tau_{:k}$  and  $t \leq k \leq h$ ,
$$
\pi_ {t} \in \arg \max  _ {\pi} H (\boldsymbol {\eta} (\tau_ {: k - 1}) + \boldsymbol {\psi} ^ {\pi} (\tau_ {: k}, \pi (\tau_ {: k}))). \tag {17}
$$
The exploration policy  $\pi_t$  is the optimal after executing the first  $t$  steps of an  $h$ -step trajectory  $\tau$ . The goal is to prove that the induction hypothesis in line (17) holds for  $t = 1$ .
Base case: The base case for  $t = h$  holds trivially, because SR does not have a dependency on the policy  $\pi$  for an  $h$ -step trajectory. Therefore the policy  $\pi$  can output any action for a trajectory sequence of length  $h$ :
$$
\max  _ {\pi} H \left(\boldsymbol {\eta} \left(\tau : h - 1\right) + \boldsymbol {\psi} ^ {\pi} \left(\tau : h, \pi \left(\tau : h\right)\right)\right) = H \left(\boldsymbol {\eta} \left(\tau : h - 1\right) + \gamma (h) \boldsymbol {e} _ {s _ {h}}\right). \tag {18}
$$
Induction Step: Suppose the induction hypothesis in line (17) holds for some  $t > 1$  and  $\pi_t$  is the maximizer of
$$
H (\boldsymbol {\eta} (\tau_ {k - 1}) + \boldsymbol {\psi} ^ {\pi_ {t}} (\tau_ {k}, \pi_ {t} (\tau_ {k}))) \tag {19}
$$
where  $t \leq k \leq h$ . For time step  $t - 1$ , we have that for some action  $a$ ,