SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} H (\boldsymbol {\eta} (\tau_ {t - 2}) + \boldsymbol {\psi} ^ {\pi_ {t}} (\tau_ {t - 1}, a)) = H (\boldsymbol {\eta} (\tau_ {t - 2}) + \gamma (t - 1) \boldsymbol {e} _ {s _ {t - 1}} + \mathbb {E} [ \boldsymbol {\psi} ^ {\pi_ {t}} (\tau_ {t}; \pi_ {t} (\tau_ {t}) | s _ {t - 1}, a ]) (20) \\ = H (\boldsymbol {\eta} (\tau_ {t - 1}) + \mathbb {E} [ \boldsymbol {\psi} ^ {\pi_ {t}} (\tau_ {t}, \pi_ {t} (\tau_ {t})) | s _ {t - 1}, a ]) (21) \\ = H \left(\mathbb {E} \left[ \boldsymbol {\eta} \left(\tau_ {: t - 1}\right) + \boldsymbol {\psi} ^ {\pi_ {t}} \left(\tau_ {: t}, \pi_ {t} \left(\tau_ {: t}\right)\right) \mid s _ {t - 1}, a \right]\right). (22) \\ \end{array}
$$
We note that the term inside the expectation is already maximized by  $\pi_t$  (by induction hypothesis). If we now set  $\pi_{t-1}$  to be equal to  $\pi_t$  for every  $t$ -step or longer trajectory and set
$$
\pi_ {t - 1} (\tau_ {: t - 1}) = \arg \max  _ {a} H (\boldsymbol {\eta} (\tau_ {: t - 2}) + \boldsymbol {\psi} ^ {\pi_ {t}} (\tau_ {: t - 1}, a)), \tag {23}
$$
then for  $t - 1\leq k\leq h$
$$
\pi_ {t - 1} \in \arg \max  _ {\pi} H (\boldsymbol {\eta} (\tau_ {: k - 1}) + \boldsymbol {\psi} ^ {\pi} (\tau_ {: k}, \pi (\tau_ {: k}))). \tag {24}
$$
B  $\eta \psi$ -Learning- Policy Gradient
Application of Q-Learning based approaches to continuous action space is not easy because finding the greedy action at any time step can be slow to be practical with large, unconstrained function approximators and nontrivial action spaces. In this work, we take a similar approach to deterministic policy gradient [55] to learn exploratory policies. The objective remains the same which is to maximize the entropy of state visitation distribution. However, it is challenging to estimate the gradient where the objective is based on the entropy term. Previous works have either used alternate optimization [47, 31] or similar objective functions [18]. The challenge is because of the expectation inside the logarithm in Equation 2. [31, 47] addressed this intractability by first estimating the visited state distribution and then using this estimate to optimize the entropy-based objective. Unfortunately, such alternating approaches are often prone to instability and slow convergence [18]. In this work, we take an alternative direction and learn a network to directly estimate the visited state distribution. The combination of predecessor trace  $\pmb{\eta}$  and successor representation  $\psi^{\pi}$  can be leveraged to estimate the state visitation distribution which is obtained using:
$$
\boldsymbol {\eta} \left(\tau_ {: T - 1}\right) + \boldsymbol {\psi} ^ {\pi} \left(\tau_ {: T}, a _ {T}\right) = \underbrace {\sum_ {t = 1} ^ {T - 1} \gamma (t) \boldsymbol {e} _ {s _ {t}}} _ {= \boldsymbol {\eta} \left(\tau_ {: T - 1}\right)} + \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} \left(\tau_ {: T}, a _ {T}\right)} \quad (\text {b y E q . 5}) \tag {25}
$$
The SR vector can be learned with gradient based optimization and provides the estimate of state visitation distribution for a given policy  $\pi$  and trajectory. The policy can utilize this estimate to learn optimal behaviors for efficient exploration in the environment.
To learn optimal behaviors for continuous action spaces,  $\eta \psi$ -Learning uses an actor-critic architecture comprising of a deterministic actor  $\pi_{\mu}(\tau)$  that provides the action and a critic to estimate the utility function. Here, both the actor and critic networks are non-Markovian and depend on the entire history of visited states. The goal of the critic network is to approximate the Q-function for a given trajectory  $\tau$  and a given action  $a \in \mathcal{A}$ . For a given trajectory  $\tau: T$ , critic computes this by combining the predecessor representation and the SR vector. The predecessor representation is fixed for a given history, implying that the critic only needs to approximate the SR  $\psi_{\theta}(\tau, a)$ . To summarize, the critic estimates the Q-function as shown below:
$$
Q _ {\theta , \exp} (\tau_ {T}, a _ {T}) = H (\boldsymbol {\eta} (\tau_ {T - 1}) + \boldsymbol {\psi} _ {\theta} (\tau_ {T}, a _ {T})), \tag {26}
$$
To update the critic network, we update the SR approximator network using temporal-difference error. The target for the SR is obtained using the action coming from the current policy  $a_{T + 1}^{\prime} = \pi_{\mu}(\tau_{:T + 1})$  and is given by
$$
\boldsymbol {y} = \boldsymbol {e} _ {s _ {T}} + \gamma (T + 1) \boldsymbol {\psi} _ {\theta} \left(\tau_ {: T + 1}, a _ {T + 1} ^ {\prime}\right). \tag {27}
$$
The SR network is updated with gradient-based learning to optimize the Mean-Squared Error between the predicted SR and the target, and the loss function  $\mathcal{L}_{SR}$  is given by
$$
\mathcal {L} _ {S R} = \left\| \boldsymbol {\psi} _ {\boldsymbol {\theta}} \left(\tau_ {: T}, a _ {T}\right) - \boldsymbol {y} \left(\tau_ {: T + 1}, a _ {T + 1} ^ {\prime}\right). \right\| ^ {2} \tag {byEq.10}
$$
Given an estimate of the SR for the current policy, we need a mechanism to update the actor network to maximize the objective. Deterministic policy gradient algorithm [55] provided a way of learning optimal policies with a deterministic actor. In this work, we formulate the gradient for the actor parameters using similar mechanism with the goal to maximize the entropy-based utility function. Proposition 2 presents a derivation of the gradients for the actor network parameters obtained by applying the chain rule on the Shannon-entropy based Q-function.
Proposition 2. Assuming the CMP satisfies [55, conditions A.1] (all functions are continuous and differentiable across all parameters) and for a  $\mu$ -parameterized policy function  $\pi_{\mu}$  the gradient with respect to  $\mu$  of the maximum entropy objective
$$
J (\pi_ {\mu}) = \mathbb {E} _ {\tau \sim \rho} [ H (\boldsymbol {\eta} (\tau_ {: - 1}) + \boldsymbol {\psi} (\tau , \pi_ {\mu} (\tau)) ]
$$
$$
\nabla_ {\mu} J (\pi_ {\mu}) = \mathbb {E} _ {\tau \sim \rho} \Big [ \sum_ {i} z _ {i} \nabla_ {\mu} \pi_ {\mu} (\tau) \nabla_ {a} \pmb {\psi} _ {i} (\tau , a) \Big | _ {a = \pi_ {\mu} (\tau)} \Big ].
$$
where  $z_{i} = -\log [\pmb{\eta}(\tau_{: - 1})_{i} + \pmb{\psi}(\tau ,\pi_{\mu}(\tau))_{i}] - 1$ ,  $H$  is the Shannon-Entropy function over the representation vectors, and the expectation over trajectories is computed with respect to some trajectory visitation distribution  $\rho$ .
Proof. We begin by rewriting the Shannon Entropy here for a  $T$ -step trajectory as
$$
H (\boldsymbol {\eta} \left(\tau_ {: T - 1}\right) + \boldsymbol {\psi} \left(\tau_ {: T}, a _ {T}\right)) = - \sum_ {i} \left(\boldsymbol {\eta} \left(\tau_ {: T - 1}\right) _ {i} + \boldsymbol {\psi} \left(\tau_ {: T}, a _ {T}\right) _ {i}\right) \log \left(\left(\boldsymbol {\eta} \left(\tau_ {: T - 1}\right) _ {i} + \boldsymbol {\psi} \left(\tau_ {: T}, a _ {T}\right) _ {i}\right), \right. \tag {28}
$$
where  $a_{T} = \pi_{\mu}(\tau_{:T})$
To simplify the notations, we will use  $\pmb{\eta}_i = \pmb{\eta}(\tau_{:T - 1})_i$  and  $\psi_{i} = \psi (\tau_{:T},a_{T})_{i}$  to represent the  $i$ th term of the predecessor and successor representation vectors. Now taking the gradient with respect to the actor parameters  $\mu$  gives:
$$
\begin{array}{l} \nabla_ {\mu} H (\boldsymbol {\eta} _ {i} + \boldsymbol {\psi} _ {i}) = - \nabla_ {\mu} \sum_ {i} \left(\boldsymbol {\eta} _ {i} + \boldsymbol {\psi} _ {i}\right) \log \left(\boldsymbol {\eta} _ {i} + \boldsymbol {\psi} _ {i}\right) (29) \\ = - \sum_ {i} \left[ \nabla_ {\mu} \left(\boldsymbol {\eta} _ {i} + \boldsymbol {\psi} _ {i}\right) \log \left(\boldsymbol {\eta} _ {i} + \boldsymbol {\psi} _ {i}\right) \right] (30) \\ = - \sum_ {i} \left[ \log \left(\boldsymbol {\eta} _ {i} + \boldsymbol {\psi} _ {i}\right) \nabla_ {\mu} \left(\boldsymbol {\eta} _ {i} + \boldsymbol {\psi} _ {i}\right) + \left(\boldsymbol {\eta} _ {i} + \boldsymbol {\psi} _ {i}\right) \nabla_ {\mu} \log \left(\boldsymbol {\eta} _ {i} + \boldsymbol {\psi} _ {i}\right) \right] (31) \\ = - \sum_ {i} \left[ \log \left(\boldsymbol {\eta} _ {i} + \boldsymbol {\psi} _ {i}\right) \nabla_ {\mu} \boldsymbol {\psi} _ {i} + \nabla_ {\mu} \boldsymbol {\psi} _ {i} \right. (32) \\ = - \sum_ {i} \left[ \log \left(\boldsymbol {\eta} _ {i} + \boldsymbol {\psi} _ {i}\right) + 1 \right] \nabla_ {\mu} \boldsymbol {\psi} _ {i} (33) \\ \end{array}
$$
Now, using the chain rule on the  $i$ th feature in SR, we obtain
$$
\nabla_ {\mu} \boldsymbol {\psi} _ {i} = \nabla_ {\mu} \pi_ {\mu} (\tau : T) \nabla_ {a} \boldsymbol {\psi} (\tau : T, a) _ {i} | _ {a = \pi_ {\mu} (\tau : T)} \tag {34}
$$
By substitution
$$
\nabla_ {\mu} H (\boldsymbol {\eta} _ {i} + \boldsymbol {\psi} _ {i}) = - \sum_ {i} [ \log (\boldsymbol {\eta} _ {i} + \boldsymbol {\psi} _ {i}) + 1 ] \nabla_ {\mu} \pi_ {\mu} (\tau ; T) \nabla_ {a} \boldsymbol {\psi} (\tau ; T, a) _ {i} | _ {a = \pi_ {\mu} (\tau ; T)}. \tag {35}
$$
Therefore, the gradient of the overall objective is
$$
\begin{array}{l} \nabla_ {\mu} J \left(\pi_ {\mu}\right) = \mathbb {E} _ {\tau \sim \rho} \left[ \nabla_ {\mu} H \left(\boldsymbol {\eta} \left(\tau : - 1\right) + \boldsymbol {\psi} \left(\tau , \pi_ {\mu} (\tau)\right) \right] \right. (36) \\ = - \mathbb {E} _ {\tau \sim \rho} \left[ \sum_ {i} [ \log (\boldsymbol {\eta} _ {i} + \boldsymbol {\psi} _ {i}) + 1 ] \nabla_ {\mu} \pi_ {\mu} (\tau ; T) \nabla_ {a} \boldsymbol {\psi} (\tau ; T, a) _ {i} | _ {a = \pi_ {\mu} (\tau ; T)} \right]. (37) \\ \end{array}
$$
Algorithm 2  $\eta \psi$ -Learning: Finite Action Space Framework
1: Initialize SR network with parameters  $\theta$  and the replay buffer  $\mathcal{B} = \{\}$   
2: Denote the predecessor feature with  $\pmb{\eta}$ , discount function with  $\gamma$ , and episode length with  $h$   
3: while Training do  
4: Collect  $\tau_{exp} = \{s_1, a_1,.., s_h\}$  using current policy  $\pi_{\theta}$  and add it to replay buffer  $\mathcal{B} = \mathcal{B} \cup \tau_{exp}$   
5: for each training step do  
6: Sample batch of  $\tau = (s_1,.., a_{l-1}, s_l) \sim \mathcal{B}$  of sequence length  $l \in \{2,.., h\}$   
7: Compute  $a' = \arg \max_{a \in A} H(\pmb{\eta}(\tau) + \pmb{\psi}_{\theta}(\tau, a))$   
8: Compute target  $\pmb{y} = \pmb{e}_{sl-1} + \gamma(l)\pmb{\psi}_{\theta}(\tau, a')$   
9: Update SR network by performing gradient step on  $||\pmb{y} - \pmb{\psi}_{\theta}(\tau; l-1, a_{l-1})||_2^2$   
10: end for  
11: end while
In Proposition 2, we derive the gradient of the actor parameters for the maximum state entropy exploration objective. Taking inspiration from algorithms [35, 17, 19] that extend Deterministic Policy Gradient (DPG) to make the optimization process stable when scaling to larger state and action space, we base our implementation to be similar to the TD3 [17] algorithm. In Appendix C, we outline the learning procedure to learn using the policy gradient derived in Proposition 2. Furthermore, we also discuss how the proposed algorithm handles continuous state spaces.
C Neural Network Architecture and Implementation
An implementation of the  $\eta \psi$ -Learning algorithm together with instructions for reproducing the experiments presented in this paper can be found at https://github.com/arnavkj1995/Eta_Psi_Learning.
$\eta \psi$ -Learning approximates the SR with a parameterized function  $\psi_{\theta}$  to learn an exploration policy and predict the state visitation distribution. Because the SR is conditioned on a trajectory  $\tau$  of variable length, we implement the function  $\psi_{\theta}$  with a Recurrent Neural Network (RNN) architecture, as outlined in Figure 6. In this architecture, the states in a trajectory  $\tau: T$  are first fed through a encoder network (E) comprising of Multi-Layer Perceptron (MLP) layers. Subsequently, the output of the Multi-Layer Perceptron (MLP) is fed through an Recurrent Neural Network (RNN) (denoted with F) architecture to compress the state sequence into one real-valued feature vector. Since, RNN are known to suffer from vanishing gradients [7], we implement the RNN with a Gated Recurrent Unit (GRU) [10]. Leveraging recurrent networks to learn the SR has been explored previously in [4, 8]. Finally, the recurrent state obtained from the RNN is concatenated with the representation of the current state and is passed through the decoder (D) with MLP layers to predict an SR vector for a given action. In the following paragraphs, we elaborate on how the proposed architecture was used to train the agent for finite and infinite action spaces.
Finite Action Space Variant For the finite action space variant, the decoder outputs a SR vector for each action  $a \in \mathcal{A}$ . This is similar to the prior method that learns Successor Features (SF) for discrete action spaces [33, 3]. Algorithm 2 describes the learning procedure for training  $\eta \psi$ -Learning to get exploratory policies.
Infinite Action Space Variant For infinite action space variant, the hidden state from the recurrent network is passed through a deterministic actor network which comprises of MLP layers. The policy network (actor) is conditioned on the hidden states because in  $\eta \psi$ -Learning the policy is a function of trajectories and not individual states. The hidden state from the recurrent network is concatenated with the action to predict the SR vector. The estimated SR vector is used to calculate the visitation distribution over one-hot embeddings of states and these SR predictions are then used to compute to loss objective for optimization. For the Reacher and Pusher tasks, we manually sub-select which dimensions of the state space are one-hot encoded. In these cases,  $\eta \psi$ -Learning learns an exploration policy that maximizes the entropy of visitation distribution across these sub-selected dimensions only. This approach to sub-selecting state dimensions is similar to prior work on maximum state entropy exploration [23, 43]. In this work, the agent is trained using similar techniques as the TD3 agent [17]. The agent keeps a single encoder and recurrent network to encode the history of
Algorithm 3  $\eta \psi$ -Learning: Infinite Action Space Framework
1: Initialize SR network with parameters  $\theta_{1},\theta_{2}$  , policy parameters  $\mu$  and the replay buffer  $\mathcal{B} = \{\}$    
2: Set target parameters equal to the main parameters:  $\theta_{targ,1} = \theta_1$ $\theta_{targ,2} = \theta_{2}$  , and  $\mu_{targ}\gets \mu$    
3: Denote the predecessor feature with  $\pmb{\eta}$  , discount function with  $\gamma$  , and episode length with  $h$    
4: while Training do   
5: Collect  $\tau_{exp} = \{s_1,a_1,\dots,s_h\}$  using target policy  $\pi_{\mu_{targ}}$  and add to replay buffer  $\mathcal{B} = \mathcal{B}\cup \tau_{exp}$    
6: for each training step  $j$  do   
7: Sample batch of  $\tau = (s_1,..a_{l - 1},s_l)\sim \mathcal{B}$  of sequence length  $l\in \{2,\dots,h\}$    
8: Compute target actions  $a^\prime = \text{clip} (\pi_{\mu_{targ}}(\tau ;l) + \text{clip} (\epsilon , - c,c),a_{Low},a_{High}),\epsilon \sim \mathcal{N}(0,1)$    
9: Compute i=arg mink∈{1,2} H(η(τ;l-1)+ ψθk(τ;l-1,a'))   
10: Compute target y=e+γ(l) ψθi(τ;l-1,a')   
11: Update the SR networks by performing gradient steps on ||y-ψθi(τ;l-1,al-1)||2, i=1,2   
12: if j % policy_update == 0 then   
13: Perform update step for policy by computing gradients using  $\sum_{i}z_{i}\nabla_{a}\psi_{\theta_{1}}(\tau ;l,a)|_{a = \pi (\tau ;l)}\nabla_{\mu}\pi_{\mu}(\tau ;l),$  where  $z_{i} = -\log (\pmb {\eta}(\tau ;l)i - \pmb{\psi}_{\theta_1}(\tau ;l,\pi_\mu (\tau ;l))_i) + 1$    
14: Update target networks with  $\theta_{targ,i}\leftarrow \rho \theta_{targ,i} + (1 - \rho)\theta_i,\qquad i = 1,2$  μtarg←ρμtarg+(1-ρ)μ   
15: end if   
16: end for   
17: end while
observed states. The encoded states are passed through two decoder networks to predict the SR vectors, which are used to represent the two critic networks. The target for SR is computed using the vector that leads to a smaller value of the two utility functions. There is a single actor network that specifies the action from the hidden state. In addition,  $\eta \psi$ -Learning-maintains a target network for all the components—encoder, recurrent, critics, and actor networks. Furthermore, similar to the TD3 algorithm, a clipped noise is added to each dimension of the action from the target network. Moreover, we also use delayed actor updates where the actor network is updated less frequently than the SR networks. Lastly, the gradients from the actor are not passed through the encoder and the recurrent networks. The procedure for training this variant is provided in Algorithm 3.
Figure 6: Network architecture to learn the SR. The states are firstly passed through an encoder (E), followed by feeding the encoded states through a RNN (GRU in our case) (F). This compresses the history of visited states, and the obtained hidden state is concatenated with the encoded state to predict the SR vector for an action using a decoder network (D).
In this work, we developed an algorithm to learn exploratory policies at convergence that can explore the state space efficiently within a finite-length trajectory. Such policies can benefit generalization in different applications like Meta-RL and episodic exploration. Maximum state entropy exploration is a potential direction for learning such policies. However, prior works are not very efficient as they either learn a Markovian policy, optimize for the state coverage over multiple long trajectories, or learn a mixture of stochastic policies. Due to these shortcomings, they are not widely used for solving tasks in RL. To address these concerns, we introduce  $\eta \psi$ -Learning and demonstrate that the proposed algorithm can learn to efficiently explore the state space within a finite length trajectory.  $\eta \psi$ -Learning achieves this by combining predecessor and successor representation to estimate the state-visitation distribution and utilizing this to optimize the entropy-based objective. [43] theoretically showed that learning such policies that achieve zero-regret is NP-Hard and we develop a practical algorithm to solve such tasks. We hope that the proposed method bridges the gap of leveraging policies learned using maximum state entropy exploration for more complex tasks in RL. As with any new approach, there are certain limitations:
- Scaling to high-dimensional inputs: Learning to explore more complex tasks with high-dimensional input spaces would require using a better representation learning method and a mechanism to estimate successor and predecessor representations. For representation learning, existing methods that use auxiliary losses, inverse/forward dynamics, or random network-based features can be used. The more challenging task is learning the SR and future works can explore leveraging methods like Successor Measures [62] or ProtoValueNetworks [15].
- Environments with changing dynamics: The learned SR depends on the environment dynamics and the policy, and we learn SR for a fixed environment in this work. However, many real-world tasks require exploration in an environment with changing dynamics (procedural environments [24, 67]). A potential direction is learning universal successor representation approximators [8] where the successor representations are conditioned on a context that defines the environment and we leave this for future research.
- Architectural priors for estimating SRs: The successor representations use an RNN which is known to suffer from vanishing gradient problems. Many real-world tasks require agents to retain information over multiple timesteps. Future research can explore having better architectural priors like Transformers or S4 that have better memory and are known to work well on complex tasks.
- Estimating predecessor representation: In this work, we computed the predecessor representation as the summation of the prior state representations. However, recent methods like Expected Eligibility Traces [64] show better sample efficiency and we leave leveraging such methods for future research.
- Computing entropy instead of state visitation distribution: While predicting future state visitations may seem harder than needed for entropy prediction, it is important to note that optimal decision depends on which states are visited multiple time steps into the future. It is possible that there exist more efficient algorithms for predicting this entropy, but to our knowledge such algorithms do not appear in the published literature, and  $\eta \psi$ -Learning is the first of this kind. Another challenge with estimating entropy is that the estimator needs to adapt to the changing policy during training. In the proposed algorithm, this is mitigated as both the entropy and policy directly depend on the estimated SR vector requiring no additional updates to estimate entropy given the policy. However, future work would involve discovering other more efficient methods for estimating the entropy induced by the future state visitations.
- Comparison to Forward-Backward Representation [63]: The Forward-Backward representations [62, 63] capture similar state visitation statistics as the predecessor and successor representation used by  $\eta \psi$ -Learning. Mathematically, the forward-backward (FB) representation factorizes the Q-function in an RL setting in a very different way than the  $\eta \psi$ -Learning algorithm. The FB method also focuses on a reward-maximization setting instead of exploration and are only conditioned on states whereas the  $\eta \psi$ -Learning algorithm conditions these representations on trajectories. During training, the focus of FB learning is on representation learning using an offline dataset where it is assumed that the agent does
not have to explore the environment. The learned FB representations are then used to solve multiple different tasks using the same representations, assuming that reward parameterization is known. However, the  $\eta \psi$ -Learning algorithm learns exploratory policies that could be used to initially explore an unknown task to determine this reward parameterization. We believe that the  $\eta \psi$ -Learning algorithm complements the FB method and integrating the two systems into a cohesive RL agent is an interesting avenue for future research.
E Discount Function
Figure 7: Illustration of the values of the entropy with different value of  $\alpha$  hyperparameter in the proposed  $\gamma$ -function for the trajectories introduced in Figure 1.
In this work, we have used a time-dependent  $\gamma$ -function. Using the gridworld example described in Figure 1, we now present how the choice of  $\gamma$ -function affects the entropy term in the objective. Suppose there are three trajectories followed by the given trace  $\tau_{:T}$ , where we denote the  $i$ -th trajectory with  $\tau^{(i)}$ . Here, Figure 1(b) shows an optimal trajectory  $(\tau^{(1)})$  which combined with the trace covers each cell of the grid with 15 steps. Figure 1(c) presents a suboptimal trajectory  $(\tau^{(2)})$  where the agent takes the right action from the current state and visits a previously observed state in the last step. Figure 1(d) shows another sub-optimal trajectory  $(\tau^{(3)})$  which takes the right action in the current state but visits the new state twice because it goes to the top right corner of the grid.