SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Arnav Kumar Jain*
Mila-Quebec AI Institute  
Université de Montréal
Lucas Lehnert†
Fundamental AI Research at Meta
Irina Rish
Glen Berseth
Animals have a developed ability to explore that aids them in important tasks such as locating food, exploring for shelter, and finding misplaced items. These exploration skills necessarily track where they have been so that they can plan for finding items with relative efficiency. Contemporary exploration algorithms often learn a less efficient exploration strategy because they either condition only on the current state or simply rely on making random open-loop exploratory moves. In this work, we propose  $\eta \psi$ -Learning, a method to learn efficient exploratory policies by conditioning on past episodic experience to make the next exploratory move. Specifically,  $\eta \psi$ -Learning learns an exploration policy that maximizes the entropy of the state visitation distribution of a single trajectory. Furthermore, we demonstrate how variants of the predecessor representation and successor representations can be combined to predict the state visitation entropy. Our experiments demonstrate the efficacy of  $\eta \psi$ -Learning to strategically explore the environment and maximize the state coverage with limited samples.
Animals and humans are very efficient at exploration compared to their data-hungry algorithms counterparts [36, 28, 50, 51, 12, 14, 65, 52]. For instance, when misplacing an item, a person will methodically search through the environment to locate the lost item. To efficiently explore, an intelligent agent must consider past interactions to decide what to explore next and avoid re-visiting previously encountered locations to find rewarding states as fast as possible. Consequently, the agent needs to reason over potentially long interaction sequences, a space that grows exponentially with the sequence length. Here, assuming that all the information the agent needs to act is contained in the current state is impossible [43].
In this paper, we present  $\eta \psi$ -Learning, an algorithm to learn an exploration policy that methodically searches through a task. This is accomplished by maximizing the entropy of the state visitation distribution of a single finite-length trajectory. This focus on optimizing the state visitation distribution of a single trajectory distinguishes our approach from prior work on exploration methods that maximize the entropy of the state visitation distribution [23, 61, 31, 41, 18]. For example Hazan et al. [23] focus on learning a Markovian policy—a decision-making strategy that is only conditioned
*Correspondence to <arnav-kumar.jain@mila.quebec>
†A part of this work was developed while Lucas Lehnert was at the Mila–Quebec AI Institute and the Université de Montréal.
(b)
Figure 1: Consider a  $4 \times 4$  gridworld for illustration. (a) The agent starts at the top left corner and takes a few actions (red arrows show the trace), (b) an optimal trajectory that covers the grid (green arrows) and can visit all the states without visiting any state twice, (c) a sub-optimal trajectory where agent visits a previously observed state in the last step and not able to visit all the states with the limited steps, (d) another sub-optimal trajectory where agent visits an observed state at an earlier step.
on the current state and does not consider which states have been explored before. A Markovian policy constrains the agent in its ability to express different exploration policies and typically results in randomizing at uncertain states to maximize the state entropy objective. While this will lead to uniformly covering the state space in the limit, such behaviors are not favorable for real-world tasks where the agent needs to maximize state coverage with limited number of interactions.
Figure 1 presents a didactic example to illustrate how an intelligent agent can learn to efficiently explore a  $4 \times 4$  grid. In this example, the agent transitions between different grid cells by selecting one of four actions: up, down, left, and right. To explore optimally, the agent would select actions that maximize the entropy of the visited state distribution of the entire trajectory. Suppose the agent started its trajectory in the top left corner of the grid (shown in Figure 1(a)) and has moved to the right twice and made one downward step (indicated by red arrows). At this point, the agent has to decide between one of the four actions to further explore the grid. For example, it could move left and follow the green trajectory as outlined in Figure 1(b). This path would be optimal in this example because every state is visited exactly once and not multiple times. However, the top action would lead to a sub-optimal trajectory as the agent would visit the previous state. To mitigate sub-optimal exploration, an intelligent agent must keep track of visited states to avoid revisiting states. Although the right action will lead to a novel state in the next step, the overall behavior will be sub-optimal as the agent will have to visit a state twice to explore the entire grid (depicted in Figures 1(c) and 1(d)). This further requires an agent to carefully plan and account for the states that would follow after taking an action.
In this work, we propose  $\eta \psi$ -Learning, an algorithm to compute an exploration strategy that methodically explores within a single finite-length trajectory—as illustrated in Figure 1(b).  $\eta \psi$ -Learning maintains two state representations: a predecessor representation [64, 2] to encode past state visitation frequencies and a Successor Representation (SR) [13] to predict future state visitation frequencies. The two representations are used to evaluate at every time step the decision that leads to covering all states as uniformly as possible. Specifically, for every potential action the agent can take, the SR is combined with the predecessor representation to predict the state visitation distribution for the current trajectory. Then, the action that results in the highest entropy of this state visitation distribution is selected for exploration. Furthermore, this exploration policy can be deterministic and does not randomize to achieve its maximum state entropy objective.
To summarize, the contributions of this work are as follows: Firstly, we propose a mechanism to combine successor [13] and predecessor [64] representations for maximizing the entropy of the state visitation distribution of a finite-length trajectory. To the best of our knowledge, this is the first work using the two representations to optimize the state visitation distribution and learn an efficient exploration policy. Secondly, we introduce  $\eta \psi$ -Learning, a method that utilizes the combination of two representations to learn deterministic and non-Markovian exploration policies for the finite-sample regime. Thirdly, we discuss how  $\eta \psi$ -Learning optimizes the entropy-based objective function for both finite and (uncountably) infinite action spaces. In Section 5 we demonstrate through empirical experiments that  $\eta \psi$ -Learning achieves optimal coverage within a single finite-length trajectory. Moreover, the visualizations presented in Section 5 demonstrate that  $\eta \psi$ -Learning learns
an exploration policy that maneuvers through the state space to efficiently explore a task while minimizing the number of times the same state is revisited.
The domain of exploration in Reinforcement Learning (RL) focuses on discovering an agent's environment via intrinsic motivation to accelerate learning optimal policies. Many existing exploration methods seek novelty by using prediction errors [45, 9, 54, 57] or pseudo-counts [58, 6, 39]. However, such methods add an intrinsic reward signal to improve sample efficiency in a single task setting. They do not explicitly learn a policy that is designed to efficiently explore a task. In contrast, we present a method for explicitly learning an efficient exploration strategy by maximizing the entropy of a single trajectory's state visitation distribution. Efficient exploration algorithms can be used to improve learning efficiency in application domains such as Meta RL [16, 70, 48, 37], Continual RL [27, 33], and Unsupervised RL [30]. For example, in Meta RL an agent needs to first explore to identify which of the previously observed tasks it is in before the agent can start exploiting rewards in the current task. In this context, VariBAD [70] maintains a belief over which task the agent is in given the observed interactions. While Zintgraf et al. argue that a Bayes-optimal policy implements an efficient exploration strategy, we propose a method that explicitly learns an efficient exploration policy, resulting in discovering rewarding states more efficiently than VariBAD (Section 5).
The core idea behind  $\eta \psi$ -Learning is the use of the predecessor and successor representations to predict the state visitation distribution induced by a non-Markovian policy for a single finite-length trajectory. Instead of using the successor representation for transfer, lifelong learning, or learning one representation that solve a set of tasks [3, 69, 4, 8, 38, 56, 22, 5, 32, 34, 1, 62], we use the successor representation to estimate the state visitation distribution and maximize its entropy. By using the successor representation in this way, the  $\eta \psi$ -Learning does not rely on density models [23, 31], an explicit transition model [61, 41], or non-parametric estimators such as k-NN [42]. In the following sections we will discuss how  $\eta \psi$ -Learning learns a deterministic exploration policy and does not rely on randomization techniques [42, 31] or mixing multiple policies to manipulate the state visitation distribution [31, 23]. Moreover, Mutti et al. [44] provide a theoretical analysis proving that efficient (zero regret) exploration is possible with a deterministic non-Markovian policy but computing such a policy is NP-hard. In this context,  $\eta \psi$ -Learning is to our knowledge the first algorithm for computing such an efficient exploration policy.
3 Maximum state entropy exploration
We formalize the exploration task as a Controlled Markov Process (CMP), a quadruple  $\mathcal{M} = \langle S, \mathcal{A}, p, \mu \rangle$  consisting of a (finite) state space  $S$ , a (finite) action space  $\mathcal{A}$ , a transition function  $p$  specifying transition probabilities with  $p(s, a, s') = \mathbb{P}(s'|s, a)$ , and a start state distribution  $\mu$  specifying probability of starting a trajectory at state  $s$  with  $\mu(s)$ . A trajectory is a sequence  $\tau = (s_1, a_1, \dots, a_{h-1}, s_h)$  of some length  $h$  that can be simulated in a CMP. A policy  $\pi$  specifies a probability distribution or density function over the action space that is sampled when selecting actions and simulating a trajectory in a CMP. This policy definition includes deterministic policies: For discrete action spaces a specific action is selected with probability one and for uncountably infinite action spaces the policy models a delta Dirac density function. In algorithms such as Q-learning [11], the policy is conditioned on the task state and specifies the probabilities of selecting a particular action. However, as illustrated in Figure 1, the past trajectory (red arrows) determines which next action leads to the best exploratory trajectory. Consequently, we consider policies that are functions of trajectories rather than just states.
The state visitation distribution of a trajectory  $\tau_{h} = (s_{1},a_{1},\dots,a_{h - 1},s_{h})$  of length  $h$  can be formally expressed in a probability vector by first encoding every state  $s_t$  as a one-hot bit vector  $\pmb{e}_{s_t}$ . Using this one-hot encoding, the  $h$ -step state visitation probability vector of the trajectory  $\tau_{h}$  can be computed by marginalizing across the time steps:
$$
\boldsymbol {\xi} _ {\gamma , \tau} = \sum_ {t = 1} ^ {h} \gamma (t) \boldsymbol {e} _ {s _ {t}}, \tag {1}
$$
where  $\gamma : \mathbb{N} \to [0,1]$  is the discount function (we denote the set of positive integers with  $\mathbb{N}$ ), such that  $\sum_{t=1}^{h} \gamma(t) = 1$ . Using the normalization in the discount function is necessary as it ensures that  $\pmb{\xi}_{\gamma,\tau}$  is a probability vector. We note that this use of a discount function is distinct from using a discount factor in common RL algorithms such as Q-learning but using a discount function is necessary as we will elaborate in the following section. The expected state visitation distribution for a policy  $\pi$  can be obtained by generating multiple trajectories using  $\pi$  and computing the average across them to get the expected state visitation distribution, denoted by  $\mathbb{E}_{\tau}[\pmb{\xi}_{\gamma,\tau}]$ .
An optimal exploration policy would achieve a similar visitation frequency for each state, as illustrated in Figure 1(b) where the optimal trajectory traverses every state once within the first 15 steps. For this trajectory the vector  $\xi_{\tau,\gamma}$  would encode a uniform probability vector, given  $\gamma(t) = \frac{1}{h}$  for any  $t$ . In fact, an optimal exploration policy  $\pi^{*}$  maximizes the entropy of this probability vector and solves the optimization problem
$$
\pi^ {*} \in \arg \max  _ {\pi} H \left(\mathbb {E} _ {\tau} \left[ \boldsymbol {\xi} _ {\gamma , \tau} \right]\right) \tag {2}
$$
where the expectation is computed across trajectories that are simulated in a CMP and follow the policy  $\pi$ . In the remainder of the paper, we will show how optimizing this objective leads to the uniform sweeping behavior illustrated in Figure 1 and the agent learns to maximize the entropy of the state visitation distribution in a single finite length trajectory. In the following section, we describe how  $\eta \psi$ -Learning optimizes the objective in Equation 2.
4  $\eta \psi$ -Learning
To learn an efficient exploration policy, we need to estimate the state visitation history and predict the distribution over future states. Consider a trajectory  $\tau = (s_1, a_1, \dots, s_{T-1}, a_{T-1}, s_T, \dots, a_{h-1}, s_h)$ . At an intermediary step  $T$ , we denote the  $T-1$ -step prefix with  $\tau_{:T-1} = (s_1, a_1, \dots, s_{T-1})$  and the suffix starting at step  $T$  with  $\tau_{T:} = (s_T, a_T, \dots, a_{h-1}, s_h)$ . Using this sub-trajectory notation, the discounted state visitation distribution in Equation 1 can be written as
$$
\boldsymbol {\xi} _ {\gamma , \tau} = \sum_ {t = 1} ^ {T - 1} \gamma (t) \boldsymbol {e} _ {s _ {t}} + \sum_ {t = T} ^ {h} \gamma (t) \boldsymbol {e} _ {s _ {t}}. \tag {3}
$$
Assuming the scenario presented in Section 3, suppose the agent has followed the trajectory  $\tau_{:T}$  until time step  $T$ . At this time step, the agent needs to decide which action  $a_{T}$  leads to covering the state space as uniformly as possible and maximizes the entropy of the state visitation distribution. The expected state visitation distribution for a policy  $\pi$  can be expressed by conditioning on the trace  $\tau_{:T}$  and a potential action  $a_{T} \in \mathcal{A}$ :
$$
\begin{array}{l} \mathbb {E} _ {\tau , \pi} \left[ \boldsymbol {\xi} _ {\gamma , \tau} \mid \tau : T, a _ {T} \right] = \mathbb {E} _ {\tau , \pi} \left[ \sum_ {t = 1} ^ {T - 1} \gamma (t) \boldsymbol {e} _ {s _ {t}} + \sum_ {t = T} ^ {h} \gamma (t) \boldsymbol {e} _ {s _ {t}} \mid \tau : T, a _ {T} \right] (4) \\ = \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})}, (5) \\ \end{array}
$$
where the vector  $\pmb{\eta}(\tau; T-1)$  is a variant of the predecessor representation [64, 2] and the vector  $\pmb{\psi}^{\pi}(\tau; T, a_T)$  is a variant of the successor representation (SR) [13]. Splitting the expected state visitation distribution into a vector  $\pmb{\eta}$  and  $\pmb{\psi}^{\pi}$  as outlined in Equation 5 is possible because we are assuming a discount function  $\gamma$  as defined in Section 3. At time step  $T$ , the two representations can be added together to estimate the expected state visitation probability vector. Simulating the proposed algorithm is analogous to effectively drawing Monte-Carlo samples from the expectation at different steps  $T$  to learn a SR and predict the expected visitation frequencies of  $\pmb{\xi}_{\gamma, \tau}$ .
The predecessor representation vector  $\pmb{\eta}(\tau_{:T - 1})$  can still be estimated incrementally similarly to the eligibility trace in TD(λ) algorithm [59] (but with a different weighting scheme that uses the discount function  $\gamma$ ). While the definition of the vector  $\pmb{\eta}(\tau_{:T})$  is similar to the definition of eligibility
3Here, we consider the Shannon entropy  $H(\pmb {p}) = -\sum_{i}\pmb{p}_{i}\log \pmb{p}_{i}$  , where the summation ranges over the entries of the probability vector  $\pmb{p}$
traces [60, Chapter 12], we do not use the predecessor trace for multi-step TD updates to learn more efficiently. Instead, the vector  $\pmb{\eta}(\tau : T - 1)$  estimates the visitation frequencies of past states—the predecessor states—to decide which states to explore next.
While the predecessor representation can be maintained using an update rule because the observed states are known, predicting future state visitation frequencies is more challenging. A potential solution is to exhaustively search through all possible sequences of trajectories starting from the current state. This is computationally infeasible and requires a dynamics model of the environment. Moreover, such a model is not always available, and learning them is prone to errors that compound for longer horizons [49, 26]. To this end, we learn a variant of the successor representation (SR), which predicts the expected frequencies of visiting future or successor states under a policy [13]. In contrast to previous methods which learn successor representation (SR) conditioned on the current state [13, 3],  $\eta \psi$ -Learning conditions the SR on the entire history of states, given by
$$
\boldsymbol {\psi} ^ {\pi} \left(\tau_ {: T}, a _ {T}\right) = \mathbb {E} _ {\tau_ {T + 1:}, \pi} \left[ \sum_ {t = T} ^ {h} \gamma (t) \boldsymbol {e} _ {s _ {t}} \mid \tau_ {: T}, a _ {T} \right]. \tag {6}
$$
Conditioning the SR on the trajectory  $\tau_{:T}$  is necessary because policy  $\pi$  is also conditioned on  $\tau_{:T}$  and therefore the visitation frequencies of future states depend on  $\tau_{:T}$ . Moreover, the expectation evaluates all possible trajectories after taking action  $a_{T}$  at time  $T$  and following policy  $\pi$  afterward. We discuss in Appendix C how the SR vectors are approximated using a recurrent neural network.
We saw in Equation 5 that the predecessor representation and successor representation can be combined to predict the state visitation distribution for a policy  $\pi$  and a trajectory-prefix  $\tau_{:T}$ .  $\eta \psi$ -Learning uses the estimated state visitation distribution to compute the entropy term in the objective defined in Equation 2. Specifically, the utility function  $Q_{\mathrm{exp}}$  approximates the entropy of the state visitation distribution for an action  $a_{T}$  at every time step. By defining
$$
Q _ {\exp} \left(\tau : T, a _ {T}\right) = H \left(\boldsymbol {\eta} \left(\tau : T - 1\right) + \boldsymbol {\psi} ^ {\pi} \left(\tau : T, a _ {T}\right)\right), \tag {7}
$$
the action that leads to the highest state visitation entropy is assigned the highest utility value. Notably, the proposed Q-function differs from prior methods using the SR, as we neither factorize the reward function [3, 4, 8, 62] nor use the SR for learning a state abstraction [32]. Optimizing the exploration Q-function stated in Equation 7 is challenging as it depends on the SR that itself depends on the policy  $\pi$  which changes during learning. Furthermore, Guo et al. [18] show that the Shannon-entropy based objective is difficult to directly optimize using gradient-based methods (due to the log term inside an expectation) [31, 47, 18]. In contrast, we outline in the following paragraphs how the entropy objective in Equation 7 can be directly optimized using either a Q-learning style method or a method based on the Deterministic Policy Gradient [55] framework for finite and infinite action spaces, respectively.
Finite action space framework Since the predecessor representation is fixed for a given trajectory  $\tau_{:T}$ , optimizing the Q-function defined in Equation 7 boils down to predicting the optimal SR for a given history  $\tau_{:T}$ . Similar to prior Successor Feature learning methods [3, 4, 33], we approximate the SR with a parameterized and differentiable function  $\psi_{\theta}$  and use a loss based on a one-step temporal difference error. Given an approximation  $\psi_{\theta}$ , the SR prediction target is obtained by the current state embedding and SR of the optimal action at the next step:
$$
\boldsymbol {y} \left(\tau : T + 1, a _ {T + 1} ^ {\prime}\right) = \boldsymbol {e} _ {s _ {T}} + \gamma (T + 1) \boldsymbol {\psi} _ {\boldsymbol {\theta}} \left(\tau : T + 1, a _ {T + 1} ^ {\prime}\right), \tag {8}
$$
where  $\tau_{:T + 1}$  is obtained by adding action  $a_{T}$  and the received next state  $s_{T + 1}$  to the trajectory  $\tau_{:T}$ . Analogous to Q-Learning, the optimal action at the next step is specified by
$$
a _ {T + 1} ^ {\prime} = \arg \max  _ {a \in \mathcal {A}} Q _ {\exp } (\tau : T + 1, a). \tag {9}
$$
Being greedy with respect to these entropy values to estimate the target leads to improving the policy  $\pi$  which in turn finds the SR for the optimal policy. (Appendix A presents a convergence analysis of this method in a dynamic programming setting.) Then, the function  $\psi_{\theta}$  is optimized using gradient descent on the loss function  $\mathcal{L}_{SR}$ , given by
$$
\mathcal {L} _ {S R} = \left| \left| \boldsymbol {\psi} _ {\boldsymbol {\theta}} \left(\tau : T, a _ {T}\right) - \boldsymbol {y} \left(\tau : T + 1, a _ {T + 1} ^ {\prime}\right) \right| \right| ^ {2}, \tag {10}
$$
where gradients are not propagated through the target  $\pmb{y}(\tau_{T+1}, a_{T+1}')$ . Finally, the optimal policy selects actions greedily with respect to the  $Q_{\mathrm{expl}}$  function. Algorithm 2 describes the training procedure for the proposed variant for finite action spaces.
Infinite action space framework Directly obtaining gradient estimates of objective defined in Equation 2 is challenging because of the expectation term in the non-linear logarithmic term. Previous approaches have used alternative optimization methods [31, 47] or resorted to a simpler noise-contrastive objective function [18]. In contrast with prior algorithms, we derived an  $\eta \psi$ -Learning variant for infinite action spaces that optimizes an actor-critic architecture using policy gradients to maximise the maximum state entropy objective. The agent uses an actor-critic architecture where actor and critic networks are conditioned on the history of visited states. The actor  $\pi_{\mu}(\tau)$  is parameterized with a parameter vector  $\mu$  and is a deterministic map from a trajectory to an action. The critic predicts the SR to estimate the utility function conditioned on a given trajectory and action. Similar to the finite action space variant, the predecessor representation is fixed for a given trajectory and the network has to predict SR  $\psi_{\theta}(\tau, a)$  for a given trajectory  $\tau$  and action  $a$ . Here, the target value of SR to update the critic is specified by the action obtained using the policy  $a_{T+1}' = \pi_{\mu}(\tau; T+1)$ , given by
$$
\boldsymbol {y} = \boldsymbol {e} _ {s _ {T}} + \gamma (T + 1) \boldsymbol {\psi} _ {\theta} \left(\tau_ {: T + 1}, a _ {T + 1} ^ {\prime}\right). \tag {11}
$$