SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In this section, we provide a detailed description of the practical implementation of  $h$ -GPI, as well as details of the domains used in our experiments. The code necessary to reproduce our results is available in the Supplemental Material.
B.1 Computing  $h$ -GPI policies with SFs and Forward-Backward Dynamic Programming
In this section, we show how to extend the Forward-Backward Dynamic Programming (FB-DP) algorithm (Efroni et al., 2020) to the discounted-SFs setting, in order to compute  $h$ -GPI policies efficiently. The algorithm is described in Algorithm 3.
The algorithm works by first computing the sets of states achievable from the current state,  $s$ , after  $t = 0, \dots, h$  time steps (line 1 of Algorithm 3). The complete procedure is described in the Forward-Pass algorithm (Algorithm 2). Next, the algorithm for computing  $h$ -GPI policies (Algorithm 3) assigns the value of states reachable in  $h$  steps using GPI bootstrapping with the SFs estimates (line 3). The values of states reachable in  $t = 1, \dots, h - 1$  time steps are then computed via dynamic programming (lines 4-5). Finally, the  $h$ -GPI action is identified in lines 6-8.
Regarding Algorithm 3, we first note that, given a state  $s$ , the algorithm's time complexity is  $O(N|\mathcal{A}||\mathcal{S}_h^{\mathrm{tot}}|)$ , where  $N$  is the maximal number of accessible states in one step (i.e., maximal "nearest neighbors" from any state), and  $|S_h^{\mathrm{tot}}|$  is the total number of reachable states in  $h$  time steps from state  $s$ .
Note that while we chose to extend the original FB-DP algorithm to design a practical procedure for computing  $h$ -GPI policies, due to its efficiency and closed-form formulation, other planning techniques could have been used (e.g., Monte Carlo Tree Search (Tesauro and Galperin, 1996; Silver et al., 2017)).
Algorithm 2: Forward-Pass
Input: Model  $\hat{m} = (\hat{p},\hat{\phi})$  ,planning horizon  $h\geq 0$  ,state s.   
1  $\mathcal{S}_0\gets \{s\} ,\mathcal{S}_t\gets \{\}\forall t\in \{1,\dots ,h\}$    
2 for  $t = 1,\ldots ,h$  do   
3 for  $s_{t - 1}\in S_{t - 1}$  do   
4  $\begin{array}{rlr}{\mathcal{S}_t\leftarrow \mathcal{S}_t\cup} & {} & \\ {\{s^{\prime}\mid \exists a\text{s.t.}\hat{p} (s^{\prime}|s_{t - 1},a) > 0\}} & {} & \end{array}$    
5 Return:  $\{\mathcal{S}_t\}_{t = 0}^h$
Algorithm 3:  $h$  -GPI with SFs and FB-DP
Input: Model  $\hat{m} = (\hat{p},\hat{\phi})$  , SFs  $\{\hat{\psi}^{\pi_i}\}_{i = 1}^n$  planning horizon  $h\geq 0$  , state  $s$  reward vector w.
1  $\{S_t\}_{t = 0}^h\gets$  ForwardPass  $(\hat{m},h,s)$
2 for  $s\in S_h$  do
3  $v_{h}(s)\gets \max_{a\in \mathcal{A}}\max_{\pi \in \Pi}\hat{\psi}^{\pi}(s,a)\cdot \mathbf{w}$
4 for  $t = h - 1,\ldots ,1$  do
5 for  $s\in \mathcal{S}_t$  do 6  $\begin{array}{rl} & v_{t}(s)\gets \max_{a\in \mathcal{A}}[\hat{\phi} (s,a)\cdot \mathbf{w} + \\ & \gamma \sum_{s^{\prime}}\hat{p} (s^{\prime}|s,a)v_{t + 1}(s)] \end{array}$
7 for  $a\in \mathcal{A}$  do
8  $(\mathcal{T}_{\hat{m}}^{*})^{h}\max_{\pi \in \Pi}\hat{q}_{\mathbf{w}}^{\pi}(s,a)\gets$ $[\hat{\phi} (s,a)\cdot \mathbf{w} + \gamma \sum_{s'}\hat{p} (s'|s,a)v_1(s)]$
9 Return: arg max $_a(T_{\hat{m}}^*)^h$  max $\pi \in \Pi$ $\hat{q}_{\mathbf{w}}^\pi(s, a)$
B.2 Environments
Below, we describe the three domains used in our experiments—all of which are depicted in Figure 6.
Figure 6: The FourRoom, Reacher, and FetchPush domains.
FourRoom. The Four Room domain (Barreto et al., 2017; Gimelfarb et al., 2021) is a gridworld with dimensions of  $13 \times 13$ , consisting of four rooms separated by walls. At each time step  $t$ , the agent occupies a cell and can move in one of four directions, denoted by the action set  $\mathcal{A} = \{\mathrm{up}, \mathrm{down}, \mathrm{left}, \mathrm{right}\}$ . Differently from related works, we introduced a  $10\%$  chance of the agent moving to a random direction, rather than the direction given by the selected action. This makes the state transition function,  $p$ , stochastic. If the destination cell happens to be a wall, the agent remains in its current cell. The grid contains three different types of objects, as illustrated in the leftmost figure in Figure 6. There are four instances of each object type distributed throughout the grid. The state space is defined as the concatenation of the agent's current x-y coordinates and a set of binary variables indicating whether each object has been picked up or not:  $S = \{0, 1, \dots, 12\}^2 \times \{0, 1\}^{12}$ . The features  $\phi(s, a, s') \in \{0, 1\}^3$  are one-hot encoded vectors that indicate the type of object present in the current cell. If there are no objects in the current cell, the corresponding features are set to zero. The goal cell which the agent has to reach is located at the upper-right of the map. It contains one instance of each object, and entering this cell results in the termination of the episode. In this domain,
a discount factor of  $\gamma = 0.95$  was used. In our experiments, we used the implementation available on MO-Gymnasium (Alegre et al., 2022b).
Reacher. The Reacher domain is a classic domain in the SFs literature (Barreto et al., 2017; Gimelfarb et al., 2021; Nemecek and Parr, 2021; Alegre et al., 2022a). It consists of a two-joint robot arm that must reach different target locations with the tip of its arm (see the middle figure in Figure 6). The agent's state space  $S \subset \mathbb{R}^6$  consists of the sine and cosine of the angles of the central and elbow joints, as well as their angular velocities. The agent's initial state is the one shown in the middle figure in Figure 6. The action space, originally continuous, is discretized using 3 bins per dimension corresponding to maximum positive torque (+1), negative torque (-1), and zero torque for each actuator. This results in a total of 9 possible actions:  $\mathcal{A} = \{-1,0, + 1\} ^2$ . Each feature  $\phi (s,a,s^{\prime}) \in \mathbb{R}^{4}$  is defined as  $\phi_{i}(s,a,s^{\prime}) = 1 - 4\Delta (\mathrm{target}_{i}), i = 1\dots 4$ , where  $\Delta (\mathrm{target}_i)$  is the Euclidean distance between the tip of the robot's arm and the  $i$ -th target's location. We used a discount factor of  $\gamma = 0.9$  in this domain. We used the implementation of this domain available on MO-Gymnasium (Alegre et al., 2022b), which is based on the MuJoco robotics simulator (Todorov et al., 2012).
FetchPush. Finally, we extended the FetchPush domain (Plappert et al., 2018), which consists of a Fetch robotic arm that must move a block to a given target position on top of a table by pushing the block with its gripper (see the rightmost figure in Figure 6). Importantly, the state space of this domain,  $S \subset \mathbb{R}^{19}$ , is high-dimensional, and its dynamics are significantly more complex than that of Reacher's. In our experiments, we removed from the states all information related to the position of targets, as these are irrelevant since targets have fixed locations. We discretized the action space similarly to how this was done when defining the Reacher domain. The reward features  $\phi(s, a, s') \in \mathbb{R}^4$  correspond to the negative distances between the block and each of the four target locations on the table. Our implementation of this domain is an adaptation of the one available in Gymnasium-Robotics (de Lazcano et al., 2023).
B.3 Parameters
In the FourRoom experiments, the SFs of each policy were learned similarly as in Alegre et al. (2022a), using Q-learning and 5 Dyna updates per time step. Each policy was trained for  $10^{6}$  time steps using a learning rate of 0.1 and epsilon-greedy exploration with a probability of selecting a random action linearly decayed from 1 to 0.05 during half of the training period.
In the Reacher and FetchPush domains, we used multi-layer perception (MLP) neural networks to learn universal successor feature approximators (USFAs) (Borsa et al., 2019). The USFAs,  $\hat{\psi}(s,a,\mathbf{w})$ , were modeled using MLPs with 4 layers of 256 neurons and ReLU non-linear activations. We used the Adam optimizer (Kingma and Ba, 2015) with a learning rate of  $3 \cdot 10^{-4}$ , and mini-batches of size 256. Recall that (as discussed in the main paper), we used vectors that form the standard basis of  $\mathbb{R}^d$  as training reward weights,  $\mathcal{M}$ , where  $d$  is the dimension of the reward features vector  $\phi$ . In our implementation, we also adopted popular DQN extensions to speed up and stabilize learning, such as prioritized experience replay (Schaul et al., 2016; Fujimoto et al., 2020) and Double Q-learning (Hasselt et al., 2016) with a target neural network updated after every 200 time steps. In both domains (Reacher and FetchPush), we trained the USFAs for a total of  $2 \cdot 10^5$  time steps. In Reacher, the epsilon-greedy exploration rate was kept fixed at 0.05, while in FetchPush, it was linearly decayed from 1 to 0.05.
We used an ensemble of  $n = 7$  probabilistic neural networks in both the Reacher and FetchPush environments to learn the model  $m_{\varphi}(S_{t + 1},\phi_t|S_t,A_t)$ . In the Reacher domain, each network was an MLP with 3 layers with 200 neurons each. In the FetchPush domain, we used an MLP with 4 layers with 400 neurons each. The probabilistic neural networks used to approximate the model were trained with early stopping based on a holdout validation subset with instances drawn from the experience buffer  $\mathcal{B}$  as commonly done when training such networks (Chua et al., 2018; Janner et al., 2019). In Figure 7, we depict the mean validation loss of the learned models both in the Reacher and FetchPush domains. As can be seen, the learned model rapidly becomes fairly accurate.
Compute. For the tabular experiments, we used an Intel i7-8700 CPU @ 3.20GHz computer with 32GB of RAM. For the experiments involving the function approximation setting, we used computers with NVIDIA A100-PCIE-40GB GPUs. For each random seed, the training period used for training
Figure 7: Mean validation loss of the learned environment model,  $\hat{m}_{\varphi}(s', \phi|s, a)$ , throughout the training phase during which  $\Pi$  is constructed (Left: Reacher domain; Right: FetchPush domain).
USFAs was of approximately 6.5 hours. We used the JAX library (Bradbury et al., 2018) for the implementation of neural networks.
C Related work
In this section, we discuss in more detail works that are related to  $h$ -GPI.
GPI and SFs. Previous works have extended GPI to safe RL (Gimelfarb et al., 2021; Feng et al., 2023), maximum-entropy RL (Hunt et al., 2019), unsupervised RL (Hansen et al., 2020), and hierarchical RL (e.g., via the options framework) (Barreto et al., 2019; Machado et al., 2023). Recently, Thakoor et al. (2022) introduced Geometric GPI (GGPI). GGPI uses geometric horizon models (GHMs) to learn the discounted future state-visitation distribution induced by particular policies to rapidly evaluate policies under a given, known reward function. The authors show that performing GPI over a particular type of non-stationary policy produces behaviors that outperform those in the agent's library policy.  $h$ -GPI, by contrast, learns a different type of model: an environment model, which is used to perform planning—i.e., action selection—rather than policy evaluation. Additionally, GGPI requires learning separate GHMs for each policy in the library, whereas  $h$ -GPI can operate with a single model that predicts the next reward features, alongside a single USFA. Secondly, GGPI assumes that the reward function is known a priori, while we exploit SFs to generalize over all linear rewards given reward features. Bagot et al. (2023) introduced GPI-Tree Search (GPI-TS), which is closely related to  $h$ -GPI. GPI-TS uses GPI bootstrapping as backup value estimates at the leaf nodes of a Monte Carlo tree search. However, GPI-TS does not employ SFs and was only employed in single-task settings. Moreover, it assumes an oracle model of the environment. We, by contrast, employ learned models to perform zero-shot transfer over multiple tasks, and we show how approximation errors in the model affect the performance of the  $h$ -GPI policy (see Theorems 1 and 2). Kim et al. (2022) introduced Constrained GPI (CGPI), which uses lower and upper bounds of the expected return of the optimal policy for a new task to constrain the action-values used when selecting the GPI policy's actions. Although CGPI is robust with respect to generalization errors in the values predicted by USFAs for new target tasks, it is sensitive to approximation errors in the values of the training tasks.  $h$ -GPI, by contrast, can inherently deal with approximation errors in the values of source/base policies. Other works have studied methods for constructing a policy library for use with GPI (Zahavy et al., 2021; Nemecek and Parr, 2021; Alver and Precup, 2022; Alegre et al., 2022a). These methods solve important orthogonal problems, and could potentially be combined with  $h$ -GPI.
Model-based RL and SFs. Russek et al. (2017) proposed a Dyna-style algorithm that performs background planning to update tabular successor representations (SRs) (Dayan, 1993). Momennejad et al. (2017) argued that this algorithm could better model human behavior on tasks with changing rewards when compared to pure model-free or model-based approaches. In this paper, by contrast, we employed online planning and leveraged GPI and SFs—both of which generalize SRs to continuous-state problems. Other works have studied methods for learning SF-based representations (Lehnert and Littman, 2020; Touati and Ollivier, 2021). Studying the behavior of  $h$ -GPI under different types of reward feature representations is an interesting direction for future research. We believe that our work
also shares similarities with the investigation performed by Tomov et al. (2021) on how model-free and model-based mechanisms are combined in human behavior.
Multi-step RL algorithms. Multi-step RL methods were extensively studied in the context of policy evaluation, both in the model-free (Hessel et al., 2018; van Hasselt et al., 2018) and the model-based settings (Yao et al., 2009; Janner et al., 2019). Model value expansion algorithms (Feinberg et al., 2018; Buckman et al., 2018; Abbas et al., 2020) are a representative example of the latter type of approach. In this paper, by contrast, we introduced a multi-step method for policy improvement, designed specifically to tackle transfer-learning settings. GX-Chen et al. (2022) introduced the  $\eta$ -return mixture, a new backup target for better credit assignment, which combines bootstrapping with standard value estimates and SFs, as a function of a parameter  $\eta$ . Efroni et al. (2018) studied multi-step greedy versions of the classic policy iteration and value iteration (Bertsekas and Tsitsiklis, 1996) dynamic programming (DP) algorithms. Efroni et al. (2020) proposed a multi-step greedy real-time DP algorithm that replaces 1-step greedy policies (used in policy improvement) with an  $h$ -step lookahead policy. The technique introduced in our paper is also related to the works of Sikchi et al. (2022) and Hansen et al. (2022), which combine planning and bootstrapping with a learned value estimate. However, unlike  $h$ -GPI, the abovementioned approaches did not address the multi-policy and zero-shot transfer settings. Another interesting use of multi-step prediction models is the work of Filos et al. (2022). In that work, the authors defined a measure of epistemic uncertainty based on the disagreement between the values resulting from applying the model-induced Bellman evaluation operator for different numbers of steps. Finally, we also emphasize to the interested reader that there exists neuroscientific evidence that multi-step planning occurs in the brains of humans and other animals (Miller and Venditto, 2021).