SlowGuess/NeurIPS_2023_OCR huggingface dataset:
For the intermediate step  $T$ , we define the discount factor for the predecessor representation for the trace as  $\gamma(t) = \frac{\alpha^{T - t}}{Z}$ , where  $\alpha$  is a scalar between  $(0,1]$ , and  $Z = \Sigma_{t=0}^{T} \alpha^{T - t} + \Sigma_{t=T}^{h} \alpha^{t - T}$  is the normalization factor. The  $\gamma$ -function for the successor representation is denoted using  $\gamma(t) = \frac{\alpha^{t - T}}{Z}$ . The proposed  $\gamma$ -function for both the representations is similar to discounting used in standard RL literature [13, 59]. Upon comparing the entropy for given trajectories with  $\alpha = 0.9, 1.0$ , we observe in Figure 7 that  $\tau^{(1)}$  being the optimal trajectory attains higher entropy when compared with  $\tau^{(2)}$  and  $\tau^{(3)}$ . The other sub-optimal trajectories  $\tau^{(2)}$  and  $\tau^{(3)}$  achieve same entropy when  $\alpha$  is set to 1.0. However, for  $\alpha = 0.9$ , the discount function  $\gamma$  emphasizes which states are visited earlier in the trajectory and assigns the lowest score to the trajectory  $\tau^{(3)}$  because this trajectory revisits states earlier in the sequence than the other options  $\tau^{(2)}$  and  $\tau^{(1)}$ . This example illustrates how the  $\gamma$ -function can be used to trade off near-term vs. long-term exploration behavior. Depending on the  $\alpha$  setting, the agent can be encouraged to avoid re-visiting states either only in the short-term or the long-term, similar to how discounting encourages maximizing short-term over long-term rewards in algorithms like Q-learning.
F Environments
For finite action space variant, we experimented with ChainMDP, RiverSwim,  $5 \times 5$  Grid-world, TwoRooms and FourRooms environments, which have finite action and state space, respectively. For infinite action space variant, we experiment with Reacher and Pusher environments where we want
Figure 8: (a) RiverSwim [58], (b) TwoRooms, and (c) FourRooms environments, respectively. The yellow block denote the initial state of the agent in each episode. The brick red regions represent the walls in the environment. When taking an action that collides with the walls, the chosen action does not change the state of the agent.
(a)
Figure 9: (a) Reacher and (b) Pusher environments for experiments with infinite action space.
the agent to move its fingertip to different locations in the environment. The environments used in this work are further described below:
ChainMDP: ChainMDP is an environment where the agent can take only move in two directions—left or right. In this work, we experiment with both deterministic and stochastic version of ChainMDP environments. The stochastic ChainMDP is similar to RiverSwim environment [58] (Figure 8(a)).
GridWorld: In the Gridworld environment (shown in Figure 1), the agent can take 4 actions to move in any of the 4 directions. In this work, we experiment with the gridworld of dimensions  $5 \times 5$ . The agent always starts in the top left corner of the grid. For the gridworld, there are multiple possible optimal trajectories, and the number of such trajectories increases exponentially with size of grid.
TwoRooms: The proposed TwoRooms environment is a gridworld with some walls. As shown in Figure 8(b), the agent starts at the center of wall between the two rooms and has to first navigate in one of the rooms, visit the starting state and then move to the other room. This makes the task challenging as the agents requires to track the trace because when the agents reaches initial state after exploring one room, the information of which room was visited should aid in going to the other room.
FourRooms: The FourRooms environment (depicted in Figure 8(c)) has four rooms connected which are connected by open shots between the walls. This task is even more challenging as the agents while navigating need to first explore the current room followed by efficiently going across different rooms.
Reacher: The Reacher environment (shown in Figure 9(a)) is a continuous control environment having a two-jointed robotic arm with continuous state and action spaces. The action space denotes the torques applied to the hinges. The state denotes the position, angles and angular velocities of the arms. The agent is tasked to maximize the entropy over the position of the fingertip.
Pusher: The Pusher environment (shown in Figure 9(b)) is a continuous control environment having a multiple-jointed robotic arm with continuous state and action spaces. The action space denotes the torques applied to the hinges. The state denotes the position, angles and angular velocities of the arms/hinges. Similar to the Reacher environment, the agent is tasked to maximize the entropy over the position of the fingertip. However, this task is harder to solve because having multiple joints leads to a larger action and state space making it a more challenging control problem.
For training and evaluation, we have used different parameters specific to each of the environment. For training, the parameters are the length of an episode and the number of episodes used for training. The number of environment steps can be obtained by multiplying these 2 parameters. For evaluation of an agent, we use different episode length for the three metrics defined to measure the performance of agents in section 5.
Name	ChainMDP	RiverSwim	Gridworld	TwoRooms	FourRooms
Training Parameters
Length of trajectory from environment	20	50	50	100	200
Number of episodes	1000	1000	1000	1000	2500
Evaluation Parameters
Horizon h for Entropy metric	20	50	50	100	200
Horizon h for State Coverage metric	20	50	50	100	200
Horizon h to measure Episode Length	100	500	200	1000	1000
Table 1: Defines the parameters of the environments with discrete actions during training and evaluation, respectively.
Name	Reacher	Pusher
Training Parameters
Length of trajectory from environment	100	200
Number of episodes	1000	1000
Evaluation Parameters
Horizon h for Entropy metric	100	200
Horizon h for State Coverage metric	100	200
Table 2: Defines the parameters of the environments with continuous actions during training and evaluation, respectively.
G Hyper Parameters
In this section, we describe the hyperparameters used for training the proposed method  $\eta \psi$ -Learning. Table 3 and Table 4 presents the list of hyperparameters for the discrete and continuous control environments, respectively. All models were trained on a single NVIDIA V100 GPU with 32 GB memory. The implementation of the proposed method was done using the RLHive [46] library.
H Meta-RL
In this work, we demonstrated how  $\eta \psi$ -Learning can learn optimal policies that can maximize the entropy of state visitation distribution. Such policies are useful for many subareas of RL where during evaluation the task is to adapt to new reward functions with minimal interactions with the environment. A challenging subproblem in such tasks is to infer the reward function. This is especially harder when the reward is sparse. Some recent works have explored adding exploratory behaviors for initial interactions with the environment to allow agent to infer the reward function. VariBAD [70] algorithm learns optimal policies that can explore during evaluation to speed up adaptation for Meta-RL tasks. VariBAD maintains a belief over the state space to explore in the environment and upon discovering the reward function adapts to the To illustrate the exploratory capabilities of  $\eta \psi$ -Learning, we compare with VariBAD as baseline in this section.
Name	Value
Batch Size	32
Sequence Length	10 / 20 / 50 / 50 / 100
α for γ-function	0.95
Encoder layers	1
Encoder output dimensions	64 / 64 / 128 / 128 / 256
Encoder activation	LeakyReLU [68]
Hidden state of GRU	64 / 64 / 128 / 128 / 256
Hidden layer dimension for SR decoder	32 / 32 / 64 / 64 / 128
Decoder activation	None
Optimizer	Adam [29]
Learning rate	3e-4
Capacity of replay buffer	200000
Table 3: Hyper parameters used for training  $\eta \psi$ -Learning. When parameters are separated by  $/. / . / . / .$  it means the corresponding hyperparameters for ChainMDP, RiverSwim, Gridworld, TwoRooms and FourRooms environments, respectively. When tuning the agent for a task, we recommend searching over  $\alpha \in \{0.9, 0.95, 0.98, 0.99\}$ , and hidden state dimension of GRU and encoder output dimensions in  $\{64, 128, 256, 512\}$ . For the replay buffer, we have used the replay buffer implemented in DreamerV2 [20], which for an episode samples a chunk of a given length.
Name	Value
Batch Size	256
Sequence Length	100
α for γ-function	0.95
Encoder layers	2
Encoder output dimensions	256
Encoder activation	LeakyReLU [68]
Hidden state of GRU	256
Hidden layer dimension for SR decoder	256
Decoder activation	None
Optimizer	Adam [29]
Learning rate	3e-4
Capacity of replay buffer	200000
Polyak constant	0.005
Grad Clip	5.0
Action noise	0.1
Target noise	0.2
Table 4: Hyper parameters used for training  $\eta \psi$ -Learning on continuous state space environments. When tuning the agent for a task, we recommend searching over  $\alpha \in \{0.9, 0.95, 0.98, 0.99\}$ , and hidden state dimension of GRU and encoder output dimensions in  $\{64, 128, 256, 512\}$ . For the replay buffer, we have used the replay buffer implemented in DreamerV2 [20], which for an episode samples a chunk of a given length.
The implementation of VariBAD provided by the authors is used to conduct this experiment. The baseline was trained for 10 million environment steps on the  $5 \times 5$  gridworld using the setup described in the paper. Two metrics are employed to compare the agents:
- State Coverage computes the fraction of the state space covered by the agent.
- Goal Search Time computes the environment steps taken to locate the sparse reward goal state. This evaluates the ability of the agent at quickly finding the reward which is essential for swift adaptation to novel tasks.
The two metrics evaluate the agents on average time taken to find the reward function and the average search completion time for covering the grid. The VariBAD [70] algorithm considered sparse reward task where a random location is sampled after each episode as the goal state and is assigned a high
reward. To compute the Goal Search Time metric, we sample a goal state randomly and record the steps taken to locate the target state. For evaluation, the metric is averaged over 16 sampled goal state for each seed. Figure 4(c) presents the comparison of VariBAD and  $\eta \psi$ -Learning on both metrics across 5 seeds. The proposed method  $\eta \psi$ -Learning achieves outperforms VariBAD on both metrics while only being trained for  $100\mathrm{K}$  environment steps. This demonstrates the efficacy of  $\eta \psi$ -Learning at exploring in environment that involves inferring the reward function during evaluation. We believe  $\eta \psi$ -Learning can be combined with a adaptation policy, where the proposed method can explore to find the reward and the adaptation policy is trained to adapt quickly to the reward function, and we leave this for future work.
I Ablation Studies
In this section, we present ablations studies to understand the gains of the proposed method  $\eta \psi$ -Learning.
I.1 MaxEnt with a recurrent network
Figure 10: Comparison of the baseline MaxEnt when trained with a recurrent network.
We conduct an experiment with a modification to the baseline MaxEnt [23] where agent observed the history of visited states. This is done to evaluate if the improvements are coming from having a recurrent policy. To this end, the state-conditioned policy in MaxEnt is replaced with a recurrent policy where the GRU [10] encodes the states observed in the trajectory. The parameters of the recurrent policy is optimized using the loss function described in MaxEnt [23]. Figure 10 presents the results of MaxEnt with a recurrent policy (named MaxEnt-GRU) where no gains are observed by having a recurrent policy and the proposed objective function used to train  $\eta \psi$ -Learning is crucial for learning optimal behaviors.
Figure 11: Comparison of  $\eta \psi$ -Learning with baseline MaxEnt when metrics are computed using multiple trajectories. MaxEnt-X denotes the evaluation with X trajectories.
I.2 Comparison across multiple trajectories
Contemporary methods on Maximum State Entropy Exploration [23, 42] were evaluated by averaging the state visitation distribution over multiple trajectories. In this work, we demonstrate that  $\eta \psi$ -Learning can achieve optimal behaviors over a single trajectory of finite length. In this study, we also explore comparison of the baseline MaxEnt when evaluated over multiple trajectories. For this evaluation, we sample a batch of trajectories and then average the state visitation distribution of trajectories. The metrics are then computed using this averaged visited state distribution. In Figure 11, we compare the Entropy and State Coverage over this averaged distribution. MaxEnt-X denotes the metric of MaxEnt after sampling X trajectories during evaluation. The proposed method  $\eta \psi$ -Learningwas evaluated using a single trajectory. We do not report metrics of  $\eta \psi$ -Learningacross multiple trajectories as we observed that the gains do not vanish with an increasing number of trajectories. The metrics for the MaxEnt algorithm improve with the increasing number of trajectories used for evaluation. With 10 or more trajectories, the baseline achieves optimal State Coverage. However,  $\eta \psi$ -Learning achieves full coverage with a single trajectory demonstrating the efficiency of the exploration policies learned using the proposed method. The baseline MaxEnt show similar behaviors by improving on the Entropy metric with more trajectories used for evaluation, whereas  $\eta \psi$ -Learning still outperforms the baseline when evaluated using a single trajectory. This demonstrates that the proposed method explores the state-space with near-equal state visitations to maximize the entropy while having optimal state coverage in a single trajectory.
I.3 Effect of the  $\alpha$  parameter
We also study the effect of the hyper-parameter  $\alpha$  of the  $\gamma$ -function (discussed in Appendix E). We note that  $\alpha$  can be selected using the same method used to select the discount factor in the standard RL. We conducted experiments with  $\alpha = \{0.8, 0.9, 0.95, 0.99\}$  across three environments- RiverSwim, TwoRooms, and FourRooms (Figure 12). On the RiverSwim environment, all methods converged with similar values across all metrics. On the TwoRooms environments, agents with  $\alpha = \{0.8, 0.99\}$  were not performing well across the three metrics. Moreover, the convergence was slower for agent with  $\alpha = 0.9$  when compared with agent trained with  $\alpha = 0.95$ . Our intuition behind this is that when  $\alpha$  is smaller, the memory of the visited states in predecessor representation ( $\eta$ ) reduces leading to a re-visitation of observed states. Whereas when  $\alpha$  is large, then the agent does similarly discount a future state at any point in the trajectory. Lastly, on the FourRooms environment, the results are similar to the TwoRooms environment but with more pronounced differences in metrics for different
Figure 12: Evaluation with different value of  $\alpha$  hyperparameter in the proposed  $\gamma$ -function.
values of  $\alpha$ . This is because FourRooms environment is harder to solve than TwoRooms environments. Notably, the Search Completion Time metric diverges for all values of  $\alpha \neq 0.95$ , which shows that  $\alpha = 0.95$  leads to optimal behavior for our task. For new tasks, we believe  $\alpha$  can be tuned similarly to the discount factor used in standard RL.
I.4 SR conditioned on state
In this section, we validate if conditioning on the history of visited states is important for learning optimal SR vectors. We consider a variant of  $\eta \psi$ -Learning where the SR is only conditioned on the current state. However, the objective proposed in Equation 7 which uses predecessor representation to compute the entropy is used for learning. We conduct experiments on continuous control tasks where we observe that conditioning on the entire trajectory sequence is important for learning optimal SR vectors (Figure 13).
I.5 Visualization of trajectories
The maneuvers taken by the learned  $\eta \psi$ -Learning agent were also visualised. For the Reacher environment, the agent was first covering the faraway states followed by covering the nearby states to the central position. Similar behaviors were observed for the Pusher environment, where the agent was moving the fingertip to different locations on the table top and tried visiting all reachable locations on the table.
J Comparison on Sparse Mountain Car
In this work, we showed how the proposed method  $\eta \psi$ -Learning can learn policies to efficiently explore in an environment. To present the applicability of the proposed method to other settings,
Figure 13: Comparison with a proposed variant where the SR is conditioned on the state only denoted as  $\eta \psi$ -Learning (State based SR). We see that conditioning on the entire trajectory is important for learning SR vectors to optimize the objective.
Figure 14: Rolled out trajectories using the learned  $\eta \psi$ -Learning agent on the Reacher (top row) and Pusher (bottom row) environments at different time steps, respectively.
we conduct an experiment on tasks with sparse rewards. We consider the Sparse MountainCar environment where the agent receives a positive reward after reaching the goal position after which the episode terminates and the agent receives no reward in other states. We chose this environment because the reward is sparse, and it is hard for the agent to discover the reward function as the agent needs to plan to explore the top of the hill. For the baselines, we compare with a random agent, TD3 [17], TD3 combined with count-based bonus [6] as an intrinsic reward (TD3-Count), and TD3 combined with first occupancy bonus [40] as an intrinsic reward (TD3-First).
We also propose a variant of our method to learn with extrinsic rewards. We modified the TD3 algorithm and call the proposed variant TD3-ηψ-Learning. In this variant, we propose to learn 2 critics- one to estimate the SR as mentioned in 4 ( $Q^{expl}$ ), and the other to estimate the returns obtained using extrinsic rewards ( $Q^{ext}$ ). The critic based on extrinsic rewards  $Q^{ext}$  is conditioned on the current state learned using the update rule defined in TD3 [17]. Notably,  $Q^{expl}$  and  $Q^{ext}$  are conditioned on the trajectory of prior visited states and current state, respectively. To learn the actor, the gradients are obtained using the overall Q-function which is defined as the linear sum of both Q-values based on extrinsic rewards and entropy-based term:  $Q = Q^{ext} + \beta Q^{expl}$ , where  $\beta$  denotes
Algorithm 4  $\eta \psi$ -Learning: TD3- $\eta \psi$ -Learning
1: Initialize SR network with parameters  $\theta_{1},\theta_{2},Q^{ext}$  network with parameters for Q-function with extrinsic rewards  $\phi_1,\phi_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,\phi_{targ,1} = \phi_1,$ $\phi_{targ,2} = \phi_2$  , and  $\mu_{targ}\gets \mu$    
3: Denote the predecessor feature with  $\eta$  , discount function with  $\gamma$  , and episode length with  $h$    
4: while Training do   
5: Collect  $\tau_{exp} = \{s_1,a_1,r_1,\dots,s_h\}$  using  $\pi_{\mu targ}$  and add it to replay buffer  $\mathcal{B} = \mathcal{B}\cup \tau_{exp}$    
6: for each training step  $j$  do   
7: Sample batch of  $\tau = (s_1,\dots ,a_{l - 1},r_{l - 1},s_l)\sim \mathcal{B}$  of sequence length  $l\in \{2,\dots,h\}$    
8: Compute target actions  $a^\prime = clip(\pi_{\mu targ}(\tau ;l) + 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 for SR  $y_{\theta} = e_{sl} + \gamma (l)\psi_{\theta_{targ,i}}(\tau_{l - 1},a^{\prime})$    
11: Update the SR networks by performing gradient steps on ||y-ψi(τ;l-1,a-l-1)||2, i=1,2   
12: Compute target for  $Q^{ext}$ $y_{\phi} = r_{l - 1} + \gamma_{ext}\min_{k\in 1,2}Q_{\phi_{targ,k}}^{ext}(sl,a')$    
13: Update the  $Q^{ext}$  networks by performing gradient steps on ||y-Qext(s1-a1)||2, i=1,2   
14: if j % policy_update == 0 then   
15: 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)+\nabla_\mu Q_{\phi_1}(s_{l-1},\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$    
16: Update target networks with  $\theta_{targ,i}\leftarrow \rho \theta_{targ,i} + (1 - \rho)\theta_i,\qquad i = 1,2$ $\mu_{targ}\gets \rho \mu_{targ} + (1 - \rho)\mu$    
17: end if   
18: end for   
19: end while
the trade-off between the two  $Q$ -functions. We have added a pseudo-code to describe the proposed algorithm in Algorithm 4.
Figure 5 presents the results on the Sparse Mountain Car environment. We compare two metrics-Return and Episode Length denoting the steps taken to reach the goal state across 5 seeds. The Average Steps metric highlights if the agent learns to solve the task with minimal interactions. We plot the mean and  $95\%$  confidence interval in shading. The proposed method outperforms TD3 and variants with count-based and first occupancy based bonuses. Through our experiment on Sparse MountainCar, we demonstrate that the proposed method can improve efficiency in standard RL tasks, especially in sparse reward environments. We leave it for future work to leverage the proposed extension in POMDP setting with more complex input spaces.