SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Changmin Yu $^{1,2}$  Neil Burgess $^{1,*}$  Maneesh Sahani $^{2,*}$  Samuel J. Gershman $^{3,*}$
$^{1}$ Institute of Cognitive Neuroscience;  $^{2}$ Gatsby Computational Neuroscience Unit; UCL, London, United Kingdom
$^{3}$ Department of Psychology, Harvard University, Cambridge, United States
* Joint senior authors
Exploration is essential in reinforcement learning, particularly in environments where external rewards are sparse. Here we focus on exploration with intrinsic rewards, where the agent transiently augments the external rewards with self-generated intrinsic rewards. Although the study of intrinsic rewards has a long history, existing methods focus on composing the intrinsic reward based on measures of future prospects of states, ignoring the information contained in the retrospective structure of transition sequences. Here we argue that the agent can utilise retrospective information to generate explorative behaviour with structure-awareness, facilitating efficient exploration based on global instead of local information. We propose Successor-Predecessor Intrinsic Exploration (SPIE), an exploration algorithm based on a novel intrinsic reward combining prospective and retrospective information. We show that SPIE yields more efficient and ethologically plausible exploratory behaviour in environments with sparse rewards and bottleneck states than competing methods. We also implement SPIE in deep reinforcement learning agents, and show that the resulting agent achieves stronger empirical performance than existing methods on sparse-reward Atari games.
The study of exploration in reinforcement learning (RL) has produced a broad range of methods [1, 2], ranging from simple methods such as pure randomization [3, 1, 4], to more sophisticated methods such as targeted exploration towards states with high uncertainty [5-7] and implicit exploration with entropy maximization [8, 9]. Intrinsic exploration, a highly effective class of methods, uses intrinsic rewards based on the agent's current knowledge of the environment, hence informing targeted exploration towards states with high predictive uncertainty or state occupancy diversity [10, 11, 6, 12]. However, existing approaches define the intrinsic reward based solely on prospective or empirical marginal information about future states, ignoring retrospective information (e.g., does a given state always precedes the goal state, hence should be more frequently traversed?). We argue that the retrospective information contains useful signals about the connectivity structure of the environment, hence could facilitate more efficient targeted exploration. For example, consider a clustered environment with bottleneck states connecting the clusters (Figure 1a), exploration based on local information (e.g., visitation counts) would discourage the agent from traversing bottleneck states, despite the key roles these states play in connecting different clusters. Guiding the agents to visit such "bottleneck" states in the face of minimal local information gain is essential in driving efficient and biologically plausible exploration. Here we study the contribution of retrospective information for global exploration with intrinsic motivation.
One of the most successful recent intrinsic exploration algorithms [12] uses the successor representation (SR; [13, 14]) to generate intrinsic rewards. The SR represents each state in terms of successor states. The row norms of the SR can be used as an intrinsic reward that generalises count-based
exploration [12]. As we discuss in Section 3, the SR contains not only prospective information, but also retrospective information about expected predecessors. This information can be utilised to construct a novel intrinsic reward which overcomes some of the problems associated with purely prospective intrinsic rewards, such as untargeted exploration, the augmented reward function is non-stationary, and asymptotic uniformity.
We provide a brief overview of background and relevant literature in Section 2, and formally introduce the novel intrinsic exploration method, Successor-Predecessor Intrinsic Exploration (SPIE), in Section 3. We propose two instantiations of SPIE for discrete and continuous state spaces, with comprehensive empirical examinations of properties of SPIE in discrete state space. We show that SPIE facilitates more efficient exploration, in terms of improved sample efficiency of learning and higher asymptotic return, through empirical evaluations on both discrete and continuous environments in Section 4.
Reinforcement Learning Preliminaries. We consider the standard RL problem in Markov Decision Processes (MDP), defined by the tuple,  $\langle S, \mathcal{A}, \mathcal{P}, \mathcal{P}_0, \mathcal{R}, \gamma \rangle$ , where  $\mathcal{S}$  is the state space,  $\mathcal{A}$  is the action space,  $\mathcal{P}: S \times \mathcal{A} \to \Delta(S)$  is the state transition distribution (where  $\Delta(S)$  is the probability simplex over  $S$ ),  $\mathcal{P}_0 \in \Delta(S)$  is the initial state distribution,  $\mathcal{R}: S \times \mathcal{A} \to \mathbb{R}$  is the reward function, and  $\gamma \in (0,1)$  is the discount factor. The goal for an RL agent is to learn the optimal policy that maximises value (expected cumulative discounted reward):  $\pi^*(a|s) = \operatorname*{argmax}_{\pi} q^\pi(s,a), \forall (s,a) \in S \times \mathcal{A}$ , where  $\pi: S \to \Delta(\mathcal{A})$  is the policy, and  $q^\pi(s,a)$  is the state-action value function:
$$
q ^ {\pi} (s, a) = \mathbb {E} _ {\mathcal {P} ^ {\pi}} \left[ \sum_ {\tau = 0} ^ {\infty} \gamma^ {\tau} \mathcal {R} \left(s _ {\tau}, a _ {\tau}\right) \mid s _ {0} = s, a _ {0} = a \right] = \mathbb {E} _ {\mathcal {P} ^ {\pi}} \left[ \mathcal {R} (s, a) + \gamma q ^ {\pi} \left(s ^ {\prime}, a ^ {\prime}\right) \right], \tag {1}
$$
where  $\mathcal{P}^{\pi}(s'|s) = \sum_{a}\pi (a|s)\mathcal{P}(s'|s,a)$  is the marginal state transition distribution given  $\pi^1$ . The second equality is the recursive form of the action value function known as the Bellman equation [15], which underlies temporal difference learning [1]:
$$
\hat {q} ^ {\pi} \left(s _ {t}, a _ {t}\right) \leftarrow \hat {q} ^ {\pi} \left(s _ {t}, a _ {t}\right) + \alpha \delta_ {t}, \quad \delta_ {t} = r _ {t} + \gamma \hat {q} ^ {\pi} \left(s _ {t + 1}, a _ {t + 1}\right) - \hat {q} ^ {\pi} \left(s _ {t}, a _ {t}\right), \tag {2}
$$
where  $\hat{q}^{\pi}(s_t,a_t)$  is the current estimate of the action values (with respect to  $\pi$ ),  $\delta_t$  is the (one-step) temporal difference (TD) error. We will study the effect of different intrinsic rewards on the performance of online TD learning (SARSA) in discrete state MDPs.
The successor representation. The SR is defined as the expected cumulative discounted future state occupancy under the policy $^2$ :
$$
\mathbf {M} [ s, s ^ {\prime} ] = \mathbb {E} _ {\mathcal {P} ^ {\pi}} \left[ \sum_ {\tau = 0} ^ {\infty} \gamma^ {\tau} \mathbb {1} \left(s _ {\tau}, s ^ {\prime}\right) | s _ {0} = s \right] = \mathbb {E} _ {\mathcal {P} ^ {\pi}} \left[ \mathbb {1} \left(s _ {0}, s ^ {\prime}\right) + \gamma \mathbf {M} \left(s _ {1}, s ^ {\prime}\right) | s _ {0} = s \right]. \tag {3}
$$
Given the recursive formulation, it is possible to learn the SR matrix online with TD learning. Given the transition tuple,  $(s_t, a_t, r_t, s_{t+1}, a_{t+1})$ , the update is
$$
\hat {\mathbf {M}} \left[ s _ {t}, s ^ {\prime} \right] \leftarrow \hat {\mathbf {M}} \left[ s _ {t}, s ^ {\prime} \right] + \alpha \delta_ {t} ^ {\mathbf {M}}, \quad \delta_ {t} ^ {\mathbf {M}} = \mathbb {1} \left(s _ {t}, s ^ {\prime}\right) + \gamma \hat {\mathbf {M}} \left[ s _ {t + 1}, s ^ {\prime} \right] - \hat {\mathbf {M}} \left[ s _ {t}, s ^ {\prime} \right], \tag {4}
$$
Note that these equations are analogous to TD learning for value function estimation, except that in this case the function being learned is a vector-valued (one-hot) representation of future states.
First-occupancy representation. The SR captures the expected cumulative discounted state occupancy over all future steps. However, in many real-world and simulated tasks, it may be preferable to reach the goal state as quickly as possible instead of as frequently as possible. In this spirit, Moskovitz et al. [16] introduced the First-occupancy Representation (FR). Formally, the FR matrix in a discrete MDP is defined by
$$
\begin{array}{l} \mathbf {F} [ s, s ^ {\prime} ] = \mathbb {E} _ {\mathcal {P} ^ {\pi}} \left[ \sum_ {\tau = 0} ^ {\infty} \gamma^ {\tau} \mathbb {1} \left(s _ {\tau} = s ^ {\prime}, s ^ {\prime} \notin \{s _ {0: \tau} \}\right) | s _ {0} = s \right] \tag {5} \\ = \mathbb {E} _ {\mathcal {P} ^ {\pi}} \left[ \mathbb {1} \left(s _ {t}, s ^ {\prime}\right) + \gamma \left(1 - \mathbb {1} \left(s _ {t}, s ^ {\prime}\right)\right) \mathbf {F} \left[ s _ {t + 1}, s ^ {\prime} \right] \mid s _ {t} = s \right], \\ \end{array}
$$
<sup>1</sup>Note that unless otherwise stated, we assume deterministic MDP, i.e.,  $\mathcal{P}(s'|s,a)$  is a delta function.
For notational simplicity, we keep the policy dependence implicit. Similar notational choice holds for all quantities discussed in the rest of the paper (F and N).
where  $\{s_{0:\tau}\} = \{s_0,s_1,\dots ,s_{\tau -1}\}$ . The recursive formulation implies that there is an efficient TD learning rule for online learning of the FR matrix. Given the transition tuple  $(s_t,a_t,r_t,s_{t + 1},a_{t + 1})$  the update rule is
$$
\hat {\mathbf {F}} [ s _ {t}, s ^ {\prime} ] \leftarrow \hat {\mathbf {F}} [ s _ {t}, s ^ {\prime} ] + \alpha \delta_ {t} ^ {\mathbf {F}}, \quad \delta_ {t} ^ {\mathbf {F}} = \mathbb {1} (s _ {t}, s ^ {\prime}) + \gamma (1 - \mathbb {1} (s _ {t}, s ^ {\prime})) \hat {\mathbf {F}} [ s _ {t + 1}, s ^ {\prime} ] - \hat {\mathbf {F}} [ s _ {t}, s ^ {\prime} ], \tag {6}
$$
Intrinsic exploration in RL. Here we focus on exploration with intrinsic motivation, where the agent augments the external rewards with self-constructed intrinsic rewards based on its current knowledge of the environment.
$$
r ^ {\operatorname {t o t}} (s, a) = r ^ {\operatorname {e x t}} (s, a) + \beta r ^ {\operatorname {i n t}} (s, a), \tag {7}
$$
where  $r^{\mathrm{ext}}(s,a)$  denotes the extrinsic environmental reward,  $r^{\mathrm{int}}(s,a)$  denotes the (possibly nonstationary) intrinsic reward, and  $\beta$  is a multiplicative scaling factor controlling the relative balance of  $r^{\mathrm{ext}}(s,a)$  and  $r^{\mathrm{int}}(s,a)$ . The intrinsic reward often operates by motivating the agent to move into under-explored parts of the state space in the absence of extrinsic reinforcement. Many types of intrinsic rewards have been proposed, including functions of state visitation counts [17-19], predictive uncertainty of value estimation [5], and predictive error of forward models [10, 6, 20]. In a closely related work, Zhang et al. [21] proposes NovelD, which constructs the episode-specific non-negative intrinsic reward based on the difference between the novelty measures of temporally adjacent states along a trajectory. However in contrast to SPIE (discussed later), the key difference is that NovelD does not explicitly utilise the retrospective information for exploration and the associated intrinsic reward is episode-dependent.
3 Successor-Predecessor Intrinsic Exploration
Existing intrinsic exploration methods construct intrinsic rewards based on either the predictive information in a temporally forward fashion (e.g., predictive error), or the empirical marginal distribution (e.g., count-based exploration). Here we argue that the retrospective information inherent in experienced trajectories, though having been largely overlooked in the literature, could also be utilised as a useful exploratory signal. Specifically, consider the environment in Figure 1a (Cluster-simple), where the discrete grid world is separated into two clusters connected by a "bottleneck" state. Whenever the starting and reward locations are in different clusters, the bottleneck state,  $s_*$ , always precedes the goal state, regardless of the trajectory taken. Hence, the frequent predecessor state (e.g.,  $s_*$ ), to the goal state should be traversed despite the fact that immediate information gain by traversing the state is minimal. In the absence of extrinsic reward, if only utilising learned prospective information based on past experience (e.g., the norm of the online-learned SR [12]), the intrinsic motivation for exploration is merely local hence would discourage transitions into bottleneck states. However, the retrospective information can be utilised to identify the state transitions that connect different sub-regions of the state space, hence incorporating the connectivity information of the state space into guiding exploration, allowing the agent to escape local exploration and navigate towards bottleneck states to reach distant regions.
We develop Successor-Predecessor Intrinsic Exploration (SPIE) algorithm utilising intrinsic rewards based on both prospective and retrospective information from past trajectories. Below we provide instantiations of SPIE based on the SR for discrete and continuous state spaces.
SPIE in discrete state space. We define the SR-Relative (SR-R) intrinsic reward, which is defined as the SR of the future state from the current state minus the sum of the SRs of the future state from all states. Formally, given a transition tuple,  $(s,a,r,s',a')$ , we define the SR-R intrinsic reward as:
$$
r _ {\mathrm {S R} - \mathrm {R}} (s, a) = \hat {\mathbf {M}} [ s, s ^ {\prime} ] - | | \hat {\mathbf {M}} [:, s ^ {\prime} ] | | _ {1} = - \sum_ {\tilde {s} \in S, \tilde {s} \neq s} \hat {\mathbf {M}} [ \tilde {s}, s ^ {\prime} ], \tag {8}
$$
The above equation holds in deterministic MDPs (i.e., when  $s'$  is a function of  $(s, a)$ ). We note that the  $j$ -th column of the SR matrix represents the expected discounted occupancy to state  $j$ , starting from every state, hence constituting a temporally backward measure of the accessibility of state  $j$  [22]. Therefore,  $r_{\mathrm{SR-R}}(s, a)$  consists of both a prospective measure  $(\hat{\mathbf{M}}[s, s'])$  and a retrospective measure  $(||\hat{\mathbf{M}}[:, s']||_1)$ , and exploring with  $r_{\mathrm{SR}}$  is an instantiation of SPIE in discrete MDPs. Intuitively,  $r_{\mathrm{SR-R}}(s)$  can be interpreted as penalising transitions leading into states  $s'$  that are frequently reached from many states other than  $s$ , hence providing an intrinsic motivation for guiding the agent towards states that are harder to reach in general, e.g., boundary states and bottleneck states. We thoroughly
investigate the individual contribution of prospective and retrospective information through ablation studies in Appendix B.4, and we observe that prospective information alone does not yield optimal exploration performance, whereas utilising only the retrospective information does not degrade exploration efficiency, indicating the importance of global topological information contained in the retrospective information for intrinsic exploration.
In a closely related work, Machado et al. [12] showed that  $r_{\mathrm{SR}}(s) = 1 / ||\hat{\mathbf{M}}[s,:]||_1$  can be used as an intrinsic reward that facilitates exploration in sparse reward settings. They additionally showed that the row norm of the online-learned SR matrix implicitly approximates the state visitation counts, so the resulting behaviour resembles count-based exploration. However, a key issue associated with  $r_{\mathrm{SR}}$  is that the asymptotic exploratory behaviour is uniformly random across all states, i.e.,  $\| \mathbf{M}[s,:] \|_1 \to 1 / (1 - \gamma)$ ,  $\forall s \in S$ . We note that exploration involves learning of both the environmental transition structure  $\mathcal{P}^{\pi}$  and the reward structure  $\mathcal{R}$ . Hence, were the SR matrix to be known a priori (hence  $\mathcal{P}^{\pi}$  could be implicitly derived), no intrinsic motivation would be introduced at any state and the resulting agent regresses back to random exploration, omitting further efficient exploration for learning  $\mathcal{R}$ . Since  $r_{\mathrm{SR-R}}$  contains the sum of columns of the SR matrix, the asymptotic uniformity in  $r_{\mathrm{SR}}$  no longer holds, yielding non-trivial intrinsic exploration even when the SR matrix is known and fixed a priori, allowing continual exploration for learning the reward structure despite sparse extrinsic reinforcement.
Analysis of  $r_{\mathrm{SR - R}}$  with pure exploration in grid-worlds. We examine exploration based on  $r_{\mathrm{SR - R}}(s)$  in discrete grid-worlds with different topologies (Fig. 1a). We first consider pure exploration in the absence of extrinsic reward, and evaluate the exploratory behaviours of 4 RL agents with different intrinsic rewards, in terms of their state coverage. The agents we consider are: vanilla SARSA [1]; SARSA with  $r_{\mathrm{SR}}$  (SARSA-SR; [12]); SARSA with  $r_{\mathrm{FR}}(s) = ||F[s, :]||_1$  (SARSA-FR; [16]); and SARSA with  $r_{\mathrm{SR - R}}$  (SARSA-SRR); the pseudocode for SARSA-SRR can be found in Appendix). We consider 4 different grid-world environments with different configurations (Figure 1a), namely,  $10 \times 10$  open-field grid (OF-small);  $10 \times 10$  grid with two rooms (Cluster-simple);  $10 \times 10$  grid with 4 rooms (Cluster-hard); and  $20 \times 20$  open-field grid (OF-large).
Exploration efficiency is quantified as the number of timesteps taken to cover  $50\%$ ,  $90\%$  and  $99\%$  of the state space. The value estimates for all states are initialised to be 0. Due to the absence of extrinsic reward, the vanilla SARSA agent is equivalent to a random walk policy, which acts as a natural baseline. We observe from Figure 1b that SARSA-SRR yields the fastest coverage of the state space amongst all considered agents. The SARSA-FR agent yields similar state coverage efficiency as SARSA. SARSA-SR performs poorly in all 4 grid-worlds, failing to achieve  $50\%$  state coverage within 8000 timesteps in all environments other than the simplest one (OF-small). Moreover, we observe that SARSA-SRR performed consistently across the 4 considered grid configurations, whereas all other agents experienced significant degradation in exploration efficiency as the size and complexity of the environments increase.
We note that in addition to improved exploration efficiency, SARSA-SRR exhibits "cycling" behaviour in pure exploration in the  $20 \times 20$  two-cluster environment (Figure 6e), spending the majority of its time exploring in one cluster and periodically traverses the "bottleneck" states to explore the opposing clusters upon sufficient coverage of the current cluster. Such "cycling" strategy exhibits short-term memory of recent states and consistent long-term planned exploration towards regions more distant in history. This is potentially advantageous for environments with non-stationary reward structures ([23]), such as real-world foraging, which require continual exploration for identifying new rewards. We verify the capability of SARSA-SRR for dealing with non-stationary reward structure in Section 4 (Figure 3).
The complexity of analysing the properties of SARSA-SRR is two-fold: the online learning of the SR matrix and the online update of the Q-values. By assuming the SR matrix is known and fixed throughout training, we observe from Figure 1c that SARSA-SRR consistently outperforms all competing methods, similar to what we observed when the SR (FR) matrix is learned online. Additionally, we observe that the exploration efficiency for all three intrinsic exploration agents drops when using the intrinsic reward constructed with the fixed SR (FR), but SARSA-SRR yields minimal decrease comparing to the significant degradation with SARSA-SR and SARSA-FR. Hence, we have empirically confirmed that the improved exploration efficiency does not stem solely from the online learning of the SR matrix, but is a property of  $r_{\mathrm{SR - R}}$ . Another long-standing issue with many existing
<sup>3</sup>We assume the policy the fixed SR matrix is dependent on is the random walk policy unless otherwise stated.
Figure 1: Evaluation of exploration efficiency in grid worlds. (a) Grid worlds with varying size and complexity. 'S' and 'G' in OF-small and Cluster-hard represents the start and goal states in the goal-oriented reinforcement learning task; colored  $G_{1}$  and  $G_{2}$  in OF-small and Cluster-hard represent the changed goal locations (see the non-stationary reward experiment in Section 4),  $s_{*}$  in Cluster-simple denote the bottleneck state. (b-c) Accumulated number of states visited against exploration timesteps, for all considered agents in all grid-worlds in with (a) online-learned SR matrix (b) and fixed SR matrix (c). All reported results are averaged over 10 random seeds (shaded area denotes mean ± 1 standard error). Hyperparameters can be found in Appendix.
intrinsic exploration methods is the non-stationary nature of the associated intrinsic bonus. By fixing the SR (FR) matrix, the associated  $r_{\mathrm{SR - R}}$  is stationary whilst still yielding high exploration efficiency, hence validating the utility of SPIE.
SPIE in continuous state space with deep RL. In order to generalise  $r_{\mathrm{SR - R}}$  to continuous state space, we replace the SR with successor features (SF; [24]).
$$
\psi^ {\pi} (s, a) = \mathbb {E} \left[ \sum_ {k = 0} ^ {\infty} \gamma^ {k} \phi_ {t + k} | s _ {t} = s, a _ {t} = a \right] = \phi \left(s _ {t + 1}\right) + \mathbb {E} \left[ \psi^ {\pi} \left(s _ {t + 1}, \pi \left(s _ {t + 1}\right)\right) | s _ {t} = s, a _ {t} = a \right] \tag {9}
$$
where  $\phi(s, a)$  is a feature representation such that  $r(s, a) = \phi(s, a) \cdot \mathbf{w}$ , with weight parameter  $\mathbf{w}$ . The recursive formulation for SF admits gradient-based learning of  $\phi$  by minimising the following squared TD loss.
$$
\delta_ {t} ^ {\mathrm {S F}} = \mathbb {E} \left[ \left(\phi \left(s _ {t}, a _ {t}\right) + \gamma \psi \left(s _ {t + 1}, a _ {t + 1}\right) - \psi \left(s _ {t}, a _ {t}\right)\right) ^ {2} \right], \tag {10}
$$
where the transition tuple  $(s_t, a_t, s_{t+1}, a_{t+1})$  can be taken from either online samples (SARSA-like) or sampled from offline trajectories (Q-learning-like). We previously noted that the column of the SR matrix provides a marginal retrospective accessibility of states, facilitating stronger exploration. However, there is no SF-analogue of the column of the SR matrix. We therefore construct the retrospective exploration objective with the Predecessor Representation (PR), which was proposed to measure how often a given state is preceded by any other state given the expected cumulative discounted preceding occupancy [25]. The formal definition for the PR matrix under discrete MDP,
$\mathbf{N} \in \mathbb{R}^{|S| \times |S|}$ , is defined as following.
$$
\mathbf {N} [ s, s ^ {\prime} ] = \mathbb {E} _ {\tilde {\mathcal {P}} ^ {\pi}} \left[ \sum_ {\tau = 0} ^ {n} \gamma^ {\tau} \mathbb {1} (s, s _ {n - \tau}) | s _ {n} = s ^ {\prime} \right] = \mathbb {E} _ {\tilde {\mathcal {P}} ^ {\pi}} \left[ \mathbb {1} (s, s _ {n}) + \gamma \mathbf {N} [ s, s _ {n - 1} ] \right], \tag {11}
$$
where the expectation is based on  $\tilde{\mathcal{P}}^{\pi}(s_t = s|s_{t + 1} = s') = \frac{\mathcal{P}^{\pi}(s,s')z(s)}{z(s')}$ , the retrospective transition model, and  $z(s) = \lim_{t\to \infty}\mathbb{E}_{\mathcal{P}^{\pi}}[\mathbb{1}(s_t = s)]$ , denotes the stationary distribution given policy  $\pi$ .
Utilising the recursive formulation for the PR matrix, we can again derive a TD-learning rule. Namely, given the transition tuple,  $(s_t,a_t,r_t,s_{t + 1},a_{t + 1})$ , we have the following update rule.
$$
\hat {\mathbf {N}} ^ {\prime} [ \tilde {s}, s _ {t + 1} ] = \hat {\mathbf {N}} [ \tilde {s}, s _ {t + 1} ] + \alpha \delta_ {t} ^ {\mathbf {N}}, \quad \delta_ {t} ^ {\mathbf {N}} = \mathbb {1} (s _ {t + 1}, \tilde {s}) + \gamma \hat {\mathbf {N}} [ \tilde {s}, s _ {t} ] - \hat {\mathbf {N}} [ \tilde {s}, s _ {t + 1} ]. \tag {12}
$$
The SR and PR have a reciprocal relationship (proof in appendix):
$$
\mathbf {N} \operatorname {d i a g} (\mathbf {z}) = \operatorname {d i a g} (\mathbf {z}) \mathbf {M}, \tag {13}
$$