Original: Learning to Prioritize Planning Updates in Model-based Reinforcement Learning
Revised: Learning to Prioritize Planning Updates in Model-based Reinforcement Learning

Original: Anonymous Author(s)

Original: AffiliationAddress

Original: Abstract
Revised: Abstract

Original: Prioritizing the states and actions from which policy improvement is performedcan improve the sample efficiency of model-based reinforcement learning systems.
Revised: Prioritizing the states and actions from which policy improvement is performed can improve the sample efficiency of model-based reinforcement learning systems.

Original: Although much is already known about prioritizing planning updates, more needsto be understood to operationalize these ideas in complex settings that involvenon-stationary and stochastic transition dynamics, large numbers of states, andscalable function approximation architectures.
Revised: Although much is already known about prioritizing planning updates, more needs to be understood to operationalize these ideas in complex settings that involve nonstationary and stochastic transition dynamics, large numbers of states, and scalable function approximation architectures.

Original: Our paper presents an online metalearning algorithm to address these needs.
Revised: Our paper presents an online meta-learning algorithm to address these needs.

Original: The algorithm finds distributions thatencode priority in their probability mass.
Revised: The algorithm finds distributions that encode priority in their probability mass.

Original: The paper evaluates the algorithm in adomain with a changing goal and with a fixed, generative transition model.
Revised: The paper evaluates the algorithm in a domain with a changing goal and with a fixed, generative transition model.

Original: Resultsshow that prioritizing planning updates from samples of the meta-learned distribution significantly improves sample efficiency over fixed baseline distributions.
Revised: Results show that prioritizing planning updates from samples of the meta-learned distribution significantly improves sample efficiency over fixed baseline distributions.

Original: Additionally, they point to a number of interesting opportunities for future research.
Revised: Additionally, they point to a number of interesting opportunities for future research.

Original:  Introduction
Revised:  Introduction

Original: Model-based reinforcement learning (MBRL) provides a computational framework for learning toachieve goals using both direct interactions with an environment, andfrom virtual interactions withan environment model.
Revised: Model-based reinforcement learning (MBRL) provides a computational framework for learning to achieve goals, using both direct interactions with an environment and virtual interactions with an environment model.

Original: Using virtual interactions to improve a policy isa process known as planning .
Revised: The process of using virtual interactions to improve a policy is known as planning .

Original: In environments where it is easier to learn a transition model than to experience transitions directly,planning can help MBRL systems rapidly learn and adapt.
Revised: Planning can help MBRL systems rapidly learn and adapt, particularly in environments where it is easier to learn a transition model than to experience transitions directly.

Original: Previous work shows that when planning under a limited computation budget, prioritizing virtualinteractions from states with large Bellman errors can lead to improved sample efficiency (Moore and
Revised: Previous work shows, that when planning under a limited computation budget, prioritizing virtual interactions from states with large Bellman errors can lead to improved sample efficiency (Moore and

Original: Atkeson, 1993).
Revised: Atkeson, 1993).

Original: This idea is used in several algorithms that apply to environments with deterministicdynamics and a small number of states (Peng and Williams, 1993; Andre et al., 1997; Wingate et al.,2005).
Revised: This idea is used in several algorithms that apply to environments with deterministic dynamics and a small number of states (Peng and Williams, 1993; Andre et al., 1997; Wingate et al., 2005).

Original: However, it remains an open question whether such planning priorities can be learned—particularly in problem settings with stochastic  dynamics and where the agent plans with a generativemodel of its environment.
Revised: Still, it remains an open question whether an RL agent can learn to prioritize the state from which it plans—particularly in problem settings with stochastic transition dynamics and where the agent maintains a generative model of its environment.

Original: Our work takes a step in this direction by showing that planning priorities can be efficiently metalearned.
Revised: Our work takes a step in this direction by showing that planning priorities can be efficiently meta learned.

Original: We present an algorithm based on the recent Bootstrap Meta Learning (Flennerhag et al. ,2022) procedure, for learning distributions of initial planning states.
Revised: We present an algorithm based on the recent Bootstrap Meta Learning (Flennerhag et al., 2022) procedure, for learning distributions of initial planning states.

Original: Experiments are performed in astochastic environment with a periodically changing goal and interactions with a generative model.
Revised: Experiments are performed in a stochastic environment with a periodically changing goal and interactions with a generative model.

Original: Results show that planning with a learned prioritization helps to focus policy improvement on statesthat matter most for planning with a limited computation budget.
Revised: Results show that planning with a learned prioritization helps to focus policy improvement on states that matter most for planning with a limited computation budget.

Original: The key contributions of this workare as follows.
Revised: The key contributions of this work are as follows

Original: • An online meta learning algorithm for prioritizing planning updates in MBRL.• Empirical evidence that  in a nonstationary domain, prioritizing updates with our learneddistribution of planning states leads to improved sample efficiency.
Revised: • An online meta learning algorithm for prioritizing planning updates in MBRL. • Empirical evidence that our proposed algorithm can improve sample efficiency in a nonstationary domain.

Original:  Reinforcement Learning Preliminaries
Revised:  Reinforcement Learning Preliminaries

Original: Our work treats the interaction between an agent and its environment as a Markov Decision Process(Sutton and Barto, 2018), defined with a set of states S and actions A , a probability distribution overstarting states p 1 , a conditional probability distribution over the next state and reward p ( S ′ , R | S, A ) ,and a discount factor γ ∈ [0 , 1) .
Revised: Our work treats the interaction between an agent and its environment as a Markov Decision Process (Sutton and Barto, 2018), defined with a set of states S and actions A , a probability distribution over starting states p 1 , a conditional probability distribution over the next state and reward p ( S ′ , R | S, A ) , and a discount factor γ ∈ [0 , 1) .

Original: The agents considered here use episodic experiences to learn a policyπ : S → P ( A ) that maximizes the expected sum of discounted rewards observed while following π :
Revised: The agents considered here use episodic experiences to learn a policy π : S → P ( A ) that maximizes the expected sum of discounted rewards observed while following π :

Original: [Equation]
Revised: [Equation]

Original: Reinforcement learning agents often use the ϵ -greedy policy, which selects uniform random actionswith probability ϵ and otherwise selects actions that maximize q ( s, a ) .
Revised: Here agents learn from experiences gathered under an ϵ -greedy policy, which selects uniform random actions with probability ϵ and otherwise selects actions that maximize q ( s, a ) .

Original: Our proposed algorithm isdesigned for general settings where the agent may not have access to the full environment state.
Revised: Our proposed algorithm is designed for general settings where the agent may not have access to the full environment state.

Original: Insuch cases, the agent constructs its own state vector x and forms an approximate value functionˆ q ( x , a ; θ )
Revised: In such cases, the agent constructs its own state vector x and forms an approximate value function ˆ q ( x , a ; θ )

Original: ≈ q ( s, a ) with a function parameterized by θ .
Revised: ≈ q ( s, a ) with a function parameterized by θ .

Original: For example, ˆ q could be a neural network.
Revised: For example, ˆ q could be a neural network.

Original: The MBRL systems studied here can be categorized under the class of Dyna algorithms (Sutton,1991).
Revised: The MBRL systems studied here can be categorized under the class of Dyna algorithms (Sutton, 1991).

Original: Dyna algorithms interleave policy improvements from both direct and virtual environmentinteractions.
Revised: Dyna algorithms interleave policy improvements from both direct and virtual environment interactions.

Original: Standard Dyna algorithms learn a value function q and an environment model m ∈49 P ( S × R ) .
Revised: Standard Dyna algorithms learn a value function q and an environment model m ∈ P ( S × R ) .

Original: However, to precisely control for the effects of different planning priorities, this workconsiders simplified Dyna algorithms that only learn value functions, and only do so through planning.
Revised: However, to precisely control for the effects of different planning priorities, this work considers simplified Dyna algorithms that only learn value functions, and only do so through planning.

Original: Expanding our results to the full class of Dyna algorithms remains an open challenge for future work.
Revised: Expanding our results to the full class of Dyna algorithms remains an open challenge for future work.

Original: Given a distribution over possible initial states d ∈ P ( S ) , the agent performs a planning update byfirst forming a transition sample (˜ x , ˜ a , ˜ x ′ , ˜ r ) from ˜ x ∼ d ( · ) , ˜ a ∼ π (˜ x ) , and ˜ x ′ , ˜
Revised: Given a distribution over possible initial states d ∈ P ( S ) , the agent performs a planning update by first forming a transition sample (˜ x , ˜ a , ˜ x ′ , ˜ r ) from ˜ x ∼ d ( · ) , ˜ a ∼ π (˜ x ) , and ˜ x ′ , ˜

Original: r ∼ m (˜ x , ˜ a ) , thenthe agent uses a learning rule such as q -learning (Watkins and Dayan, 1992) to update its valueparameters θ .
Revised: r ∼ m (˜ x , ˜ a ) , then the agent uses a learning rule such as q -learning (Watkins and Dayan, 1992) to update its value parameters θ .

Original: Notice that d determines the initial state from which planning occurs and, therefore, itsprobability mass implicitly encodes the priority of planning updates.
Revised: Notice that d determines the initial state from which planning occurs and, therefore, its probability mass implicitly encodes the priority of planning updates.

Revised: States with high mass are likely to be updated more than states with lower mass.

Original: In what follows we propose ameta learning algorithm for learning a distribution d whose probabilities are parameterized by  η .
Revised: In what follows we propose a meta learning algorithm for learning a distribution d whose probabilities are parameterized by the vector η .

Original:  Meta Learning a Priority over Initial Planning States
Revised:  Meta Learning a Priority over Initial Planning States

Original: This work is guided by two essential questions.
Revised: This work is guided by two essential questions.

Original: What is a good prioritization of states for planning?
Revised: What is a good prioritization of states for planning?

Original: And, can such planning priorities be learned?
Revised: And, can such planning priorities be learned?

Original: As noted earlier, prioritized sweeping (Moore and Atkeson, 1993) provides a suggestive example of a useful prioritization.
Revised: As noted earlier, work by Moore and Atkeson (1993) provides a suggestive example of a useful prioritization.

Original: Prioritized sweeping uses Bellman error as a heuristic which determines planning priority and improves sample efficiency.
Revised: Moore and Atkeson’s approach uses Bellman error as a heuristic to determine planning priority. However, this strategy remains fixed while the agent performs policy improvement.

Original: Ourapproach attempts to answer both questions simultaneously by leveraging recent developments inmeta-learning.
Revised: Our approach attempts to answer both questions simultaneously by leveraging recent developments in meta-learning.

Original: We propose an algorithm which prioritizes states according to their contribution toreducing parameter error.
Revised: We propose an algorithm that prioritizes states according to their contribution to reducing parameter error in θ .

Original: Our algorithm,  Adaptive Prioritization (AP), meta-learns the distributiond online using the reduction in parameter error to guide the distribution’s shape.
Revised: Our algorithm, which we term Adaptive Prioritization (AP), meta-learns the distribution d online using the reduction in parameter error to guide the distribution’s geometry.

Original: The procedure isdescribed in full in the Appendix, including pseudocode (see Algorithm 1).
Revised: The procedure is described in full in the Appendix, including pseudocode (see Algorithm 1).

Original: Intuitively, our meta loss is designed to produce a distribution which prioritizes states leadingto the largest reductions in parameter error.
Revised: Intuitively, our meta loss is designed to prefer distributions whose encoded priorities lead to large reductions in parameter error over θ .

Original: This can be achieved by minimizing the squared Euclidean difference between some  target parameters θ ∗ and theexpected value parameters afterperforming a planning update starting from θ , using step size α and learning update ∆(˜ x , θ ) ,¯ θ ( η ) ≜ θ + α (cid:80) ˜
Revised: We achieve this by minimizing the squared Euclidean difference between some vector of target parameters θ ∗ and the parameters obtained after a planning update is applied to the current θ . Although this idea is quite general, one must ensure the updated parameters, denoted ¯ θ are a differentiable function of the meta parameters η . Here we update θ with an expected planning update 1 using step size α and learning rule ∆(˜ x , θ ) ,

Original: x , ˜ a π (˜ a | ˜ x )  d (˜ x ; η)∆(˜ x , θ ) .
Revised: [Equation] Here, meta gradients will propagate through ¯ θ ( η ) directly into the distribution model d (˜ x ; η ) .

Revised: Other ways of imposing differentiability could involve the use of score functions (Kleijnen and Rubinstein, 1996), or the judicious placement of stop-gradients (Bengio et al., 2013).

Original: Here, our learning update  is a semi-gradient q -learningstep.
Revised: Our learning update ∆(˜ x , θ ) is chosen to be the semi-gradient q -learning rule.

Original: However, it could also apply to other gradient-based learning rules.
Revised: Though, we believe the loss is amenable to other gradient-based learning rules such as the Deep Q-Network rule (Mnih et al., 2015), and its many derivatives (Van Hasselt et al., 2016; Bellemare et al., 2017).

Original: To compute the targetparameters θ ∗ , we use the Bootstrapped Meta Learning procedure (Flennerhag et al., 2022).
Revised: To compute the target parameters θ ∗ , we use the Bootstrapped Meta Learning procedure (Flennerhag et al., 2022).

Original: Thisprocedure uses the meta-learner’s own output to approximate θ ∗
Revised: This procedure uses the meta-learner’s own output to approximate θ ∗

Original: ≈ ˆ θ .
Revised: ≈ ˆ θ .

Original: In the MBRL setting, thebootstrapped targets are produced by  performing l  updates  of the form ˆ θ ≜ ˆ θ + α ∆(˜ x , ˆ θ  ) where thefirst such update is performed on top of ¯ θ .
Revised: In the MBRL setting, the bootstrapped targets are produced by first computing ¯ θ then performing l additional planning updates on top of ¯ θ with samples from the agent’s transition model. Because these additional updates are sample-based, we drop the bar notation and use θ ( l ) to mean the parameters after l planning updates were applied to ¯ θ .

Original: A final update is applied using  the  direct interaction from the environment ( x , a , x ′ , r ) , such that ˆ θ ≜ˆ θ  + α ∆( x ,ˆ θ  ) .
Revised: A final update is applied using experience from the most recent direct interaction with the environment ( x , a , x ′ , r ) , such that ˆ θ ≜ θ ( l ) + α ∆( x , θ ( l ) ) .

Revised: This grounds the updates to ensure progress is made toward an optimum θ ∗ .

Original: Here we suppress the dependence  on ηbecause the meta loss applies a stop-gradient to this term:
Revised: Our meta loss is given by the following equation, where we suppress the dependence of ˆ θ on η because the meta loss applies a stop-gradient to this term:

Original: [Equation]
Revised: [Equation]

Original:  Our algorithm minimizes meta loss (2) with online samples of the direct environment interactionsusing the Adam optimizer (Kingma and Ba, 2014), which enables thepriority distribution d to track toa changing environment. Intuitively, a distribution d would be preferred by the meta-loss if, under thesame number of planning steps, it brought the parameters closer to optimality.
Revised: Measuring learning progress in parameter space captures the heuristic priority of Bellman error which Moore and Atkeson (1993) and many others found effective, but it additionally allows for the heuristic to be scaled to settings where the agent uses function approximation. Our algorithm minimizes meta loss (3) with online samples of the direct environment interactions using the Adam optimizer (Kingma and Ba, 2014). The use of online updates enables the distribution d to track its priorities to changes in the environment.

Original: The learned prioritydistributions will always place mass on states with non-zero parameter error.
Revised: Furthermore, the learned distributions always place mass on states with non-zero parameter error.

Original: This suggests that undersome mild realizability assumptions, convergence to an optimal value function is possible (Li and Littman, 2008).
Revised: This suggests that under some mild realizability assumptions, convergence to an optimal value function can be guaranteed under the analysis of Li and Littman (2008), who view the problem as an instance of asynchronous value iteration.

Original:  Empirical Results
Revised:  Empirical Results

Original: We evaluate the Adaptive Prioritization algorithm in the TMaze; an episodic grid-world environmentpictured in Figure 1a.
Revised: We evaluate the Adaptive Prioritization algorithm in the TMaze; an episodic grid-world environment pictured in Figure 1a.

Original: The TMaze is a non-stationary domain in which algorithms capable of adaptingto a changing reward structure stand to perform well.
Revised: The TMaze is a non-stationary domain in which algorithms capable of adapting to a changing reward structure stand to perform well.

Original: Our experiments demonstrate that in thisdomain, a learned prioritization scheme can outperform fixed distributions, even those which arehand-tailored to the domain to help the agent learn efficiently.
Revised: Our experiments demonstrate that in this domain, a learned prioritization scheme can outperform fixed distributions, even those which are hand-tailored to the domain to help the agent learn efficiently.

Original: In the TMaze, an agent begins at a starting state and must navigate a vertical hallway, then turn left orright at a junction.
Revised: In the TMaze, an agent begins at a starting state and must navigate a vertical hallway, then turn left or right at a junction.

Original: Reaching a state at either the left or right of the horizontal hallway results in thetermination of an episode.
Revised: Reaching a state at either the left or right of the horizontal hallway results in the termination of an episode.

Original: One of the terminal states emits a reward of +1 while the other emits 0 .
Revised: One of the terminal states emits a reward of +1 while the other emits 0 .

Original: Every 600 episodes the rewards are swapped between terminal states.
Revised: Every 600 episodes the rewards are swapped between terminal states.

Original: From the agent’s perspective,the TMaze is thus non-Markov and non-stationary.
Revised: From the agent’s perspective, the TMaze is thus non-Markov and non-stationary.

Original: At any timestep a random transition may occurwith probability ϵ env .
Revised: At any timestep a random transition may occur with probability ϵ env .

Original: A key element of the TMaze is that under the optimal policy only the values ofcertain states change.
Revised: A key element of the TMaze is that under the optimal policy only the values of certain states change.

Original: The values of states along the vertical hallway do not change when the rewardis swapped, while the values of states in the horizontal hallway do change.
Revised: The values of states along the vertical hallway do not change when the reward is swapped, while the values of states in the horizontal hallway do change.

Original:  Our experiments study agents learning to act in the TMaze environment.
Revised: [Figure] Our experiments study agents learning to act in the TMaze environment.

Original: All agents behave under anϵ -greedy policy while attempting to learn optimal q -values.
Revised: All agents behave under an ϵ -greedy policy while attempting to learn optimal q -values.

Original: In all our experiments, we use a tabularrepresentation of states and actions while agents learn values which are a linear function of states andactions.
Revised: In all our experiments, we use a tabular representation of states and actions while agents learn values which are a linear function of states and actions.

Original: We consider several non-adaptive baseline learning systems.
Revised: We consider several non-adaptive baseline learning systems.

Original: The first baseline is intended to establishthe supremacy of model-based algorithms in our domain; it is a Q-Learning agent, which learnswithout a model, directly from interaction with the environment.
Revised: The first baseline is intended to establish the supremacy of model-based algorithms in our domain; it is a Q-Learning agent, which learns without a model, directly from interaction with the environment.

Original: We then consider two MBRL agentswith fixed priority distributions: Uniform, and Domain-Specific.
Revised: We then consider two MBRL agents with fixed priority distributions: Uniform, and Domain-Specific.

Original: Uniform places uniform probabilityover all non-terminal states.
Revised: Uniform places uniform probability over all non-terminal states.

Original: The domain-specific distribution is based on the observation that the[Figure] value of states in the vertical portion of the TMaze will not change when the reward switches betweenterminal states.
Revised: The domain-specific distribution is based on the observation that the value of states in the vertical portion of the TMaze will not change when the reward switches between terminal states.

Original: Thus, the agent has less need to plan from these states.
Revised: Thus, the agent has less need to plan from these states.

Original: Depictions of the uniform anddomain-specific distributions are shown in Figures 1b and 1c.
Revised: Depictions of the uniform and domain-specific distributions are shown in Figures 1b and 1c.

Original: Figure 2a shows the total rewards accumulated during training for each agent, while Figure 2b showsthe average reward per timestep over the course of the experiment.
Revised: Figure 2a shows the total rewards accumulated during training for each agent, while Figure 2b shows the average reward per timestep over the course of the experiment.

Original: All agents experienced 100,000timesteps of interaction with the environment.
Revised: All agents experienced 100, timesteps of interaction with the environment.

Original: We observe that the adaptive agent achieves thegreatest total reward in this setting by a statistically significant margin.
Revised: We observe that the adaptive agent achieves the greatest total reward in this setting by a statistically significant margin.

Original: Further, we observed thatthe adaptive agent’s average reward increases over the course of training.
Revised: Further, we observed that the adaptive agent’s average reward increases over the course of training.

Original: While that of uniform anddomain-specific plateau.
Revised: While that of uniform and domain-specific plateau.

Original: Even though the adaptive and baseline algorithms use the same amount ofplanning, and the same amount of interaction with the environment, the search control distributionlearned by the adaptive algorithm allows planning steps to much more effectively update the agentstowards a useful value function for control.  
Revised: Even though the adaptive and baseline algorithms use the same amount of planning, and the same amount of interaction with the environment, the prioritization distribution learned by the adaptive algorithm allows planning steps to much more effectively update the agent towards a useful value function for control.  Related Work

Revised: Sutton and Barto (2018) summarize the problem of selecting an initial state and action for the simulation of virtual experiences generated by the model as search control .

Original: Conclusion
Revised: Despite the significant impact search control can have on sample efficiency, the problem has not received a great deal of attention compared to other aspects of RL research. However in one recent study, Pan et al. (2020) examine the effectiveness of a new priority heuristic—one that prioritizes regions of the state space where it is “difficult" to approximate the value. This is distinct from other work that shows planning should be avoided in regions where the model is a poor approximation to the environment’s observation process (Abbas et al., 2020), such as in offline RL (Buckman et al., 2020).

Revised: The importance of prioritizing some states over others throughout learning has connections to experience replay (Lin, 1992).

Revised: Schaul et al.

Revised: (2015) showed how to bias replay samples with the temporal difference error of recently encountered transitions, and how this can lead to significant improvements in learning performance.

Revised: But as Van Hasselt et al.

Revised: (2019) points out, a replay buffer is nothing more than a non-parametric distribution model of transitions.

Revised: Thus they remain limited as a representation of priority distributions, because they are unable to generalize priority between related experiences.

Revised:  Conclusion

Original: In this work we introduced the Adaptive Prioritization algorithm and demonstrated the potentialmeta-learning offers to adaptively adjust the planning procedure.
Revised: In this work we introduced the Adaptive Prioritization algorithm and demonstrated the potential meta-learning offers to adaptively adjust the planning procedure.

Original: Ultimately, we found that the
Revised: Ultimately, we found that the

Original: Adaptive Prioritization algorithm’s ability to shape its planning distribution continually allowed thisalgorithm to achieve superior performance in the non-stationary, episodic TMaze task.
Revised: Adaptive Prioritization algorithm’s ability to shape its planning distribution continually allowed this algorithm to achieve superior performance in the non-stationary, episodic TMaze task.

Original: We believethis result signals a fruitful use for meta-learning techniques to improve planning algorithms beyondheuristic based methods to initialize planning.
Revised: We believe this result signals a fruitful use for meta-learning techniques to improve planning algorithms beyond heuristic based methods to initialize planning.

Original: In future work, we intend to expand our experiments tosettings where a forward transition model is learned simultaneously with the prioritization distribution,and use more powerful function approximators to learn useful distributions in even more complexenvironments.
Revised: In future work, we intend to expand our experiments to settings where a forward transition model is learned simultaneously with the prioritization distribution, and use more powerful function approximators to learn useful distributions in even more complex environments.

Original: Eventually, we hope to incorporate adaptation through meta-learning in other aspectsof planning algorithms.
Revised: Eventually, we hope to incorporate adaptation through meta-learning in other aspects of planning algorithms.

Original: A Appendix 
Revised: A Appendix A.1 Adaptive Planning Initialization Pseudocode

Original: A.1 Adaptive Planning Initialization Pseudocode

Original: We provide a detailed description of the Adaptive Prioritization algorithm during each step ofinteraction with the environment.
Revised: We provide a detailed description of the Adaptive Prioritization algorithm during each step of interaction with the environment.

Original: Given the most recent observation x , AP selects an action aaccording to an ϵ -greedy policy derived from its current action-value function.
Revised: Given the most recent observation x , AP selects an action a according to an ϵ -greedy policy derived from its current action-value function.

Original: The agent takes theselected action in the environment and receives a reward r and observation x ′ .
Revised: The agent takes the selected action in the environment and receives a reward r and observation x ′ .

Original: Following this, theremainder of the algorithm can be broken down into three steps.
Revised: Following this, the remainder of the algorithm can be broken down into three steps.

Original: First, the agent performs k planningupdates using the  distribution d .
Revised: First, the agent performs k planning updates using the learned prioritization distribution d .

Original: During each update, the agent samples a state using d and selects anaction according to its current ϵ -greedy policy.
Revised: During each update, the agent samples a state from d and selects an action according to its current ϵ -greedy policy.

Original: Given the sampled state and selected action, the agentreceives a next state and reward from an oracle model of the environment.
Revised: Given the sampled state and selected action, the agent receives a next state and reward from its transition model m .

Original: This produces a tuple ofvirtual experience ( ˜ x , ˜ a , ˜ x ′ , ˜ r ) .
Revised: This produces a tuple of virtual experience ( ˜ x , ˜ a , ˜ x ′ , ˜ r ) .

Original: For each such tuple the agent updates θ by taking a semi-gradientq-learning step.
Revised: For each such tuple the agent updates θ by taking a semi-gradient q-learning step.

Revised: Lines 4-8 of Algorithm 1 reflect this planning procedure.

Original: After planning, the agent’s updated parameters θ are used to generate the parameter vector ¯ θ accordingto the ExpectedUpdate subroutine (Algorithm 3).
Revised: After planning, the agent’s updated parameters θ are used to generate the parameter vector ¯ θ according to the ExpectedUpdate subroutine (Algorithm 2).

Original: ExpectedUpdate calcuates the change in parametervalues which would occur were a planning update to occur from each possible state-action pair.
Revised: ExpectedUpdate calcuates the change in parameter values which would occur were a planning update taken from each possible state-action pair.

Original: Eachchange in parameter values is weighted by the probability of the state being sampled from d andthe action being selected by the current ϵ -greedy policy.
Revised: Each change in parameter values is weighted by the probability of the state being sampled from d and the action being selected by the current ϵ -greedy policy, as shown in Equation 2.

Original: The result is a weighted average update.
Revised: The update is thus the expected change in parameter values.

Original: This update is applied to θ to yield ¯ θ .
Revised: This update is applied to θ to yield ¯ θ .

Original: This expected update results in each component of η having anon-zero gradient when differentiating the meta-objective.
Revised: This expected update results in each component of η having a non-zero gradient when differentiating the meta-objective.

Original: Next, the BootstrapTarget subroutine (Algorithm 2) receives ¯ θ as input and produces a set of targetparameters ˆ θ .
Revised: Next, the BootstrapTarget subroutine (Algorithm 3) receives ¯ θ as input and produces a set of target parameters ˆ θ .

Original: Ideally, the meta-loss would compute the difference between ¯ θ and the optimalparameters θ ∗ .
Revised: Ideally, the meta-loss would compute the difference between ¯ θ and the optimal parameters θ ∗ .

Original: Of course, optimal parameters are not available to our agent.
Revised: Of course, optimal parameters are not available to our agent.

Original: Instead, we use theidea of bootstrapping (Flennerhag et al., 2022).
Revised: Instead, we use the idea of bootstrapping (Flennerhag et al., 2022).

Original: BootstrapTarget performs ℓ additional planningsteps on top of ¯ θ by sampling starting states from d .
Revised: BootstrapTarget performs ℓ additional planning steps on top of ¯ θ by sampling starting states from d .

Original: Finally, an update is made using the tupleveridical experience ( x , a , x ′ , r ) collected through direct interaction with the environment.
Revised: Finally, an update is made using the tuple veridical experience ( x , a , x ′ , r ) collected through direct interaction with the environment.

Original: Theresulting parameter vector ˆ θ approximates θ ∗ .
Revised: The resulting parameter vector ˆ θ approximates θ ∗ .

Original: Adaptive Prioritization optimizes to reduce theparameter difference between ˆ θ and ¯ θ .
Revised: Adaptive Prioritization optimizes to reduce the parameter difference between ˆ θ and ¯ θ .

Original: A stop gradient prevents any adjustment to η through thetarget parameters.
Revised: A stop gradient prevents any adjustment to η through the target parameters.

Original: Only ¯ θ is differentiated through, and so the distribution is shifted in favour ofstates which have the greatest effect on the parameter difference during the ExpectedUpdate step.
Revised: Only ¯ θ is differentiated through, and so the distribution is shifted in favour of states which have the greatest effect on the parameter difference during the ExpectedUpdate step.

Original: 1: Receive x 1 from environment.
Revised: 1: Receive x 1 from environment.

Original: 2: for t = 1 , 2 , 3 , · · · do 3:
Revised: 2: for t = 1 , 2 , 3 , · · · do 3:

Original: Take ϵ -greedy a t from x t and receive x t +1 and r t +1 .
Revised: Take ϵ -greedy a t from x t and receive x t +1 and r t +1 .

Original: 4: # Perform planning update.
Revised: 4: # Perform planning update.

Original: 5: for 1 , · · · , k do 6:
Revised: 5: for 1 , · · · , k do 6:

Original: Take ϵ -greedy ˜ a from ˜ x ∼ d ( η ) .
Revised: Take ϵ -greedy ˜ a from ˜ x ∼ d ( η ) .

Original: 7: ˜ x ′ , ˜
Revised: 7: ˜ x ′ , ˜

Original: r ∼ m (˜ x , ˜ a ) .
Revised: r ∼ m (˜ x , ˜ a ) .

Original: 8: θ ← θ + α [˜ r + γ max ˜ a ′ ˆ q (˜ x ′ , ˜ a ′ ; θ ) − ˆ q (˜ x , ˜ a ; θ )] ∇ θ ˆ q (˜ x , ˜ a ; θ ) 9: # Get expected and target parameters, then compute meta-loss.
Revised: 8: θ ← θ + α [˜ r + γ max ˜ a ′ ˆ q (˜ x ′ , ˜ a ′ ; θ ) − ˆ q (˜ x , ˜ a ; θ )] ∇ θ ˆ q (˜ x , ˜ a ; θ ) 9: # Get expected and target parameters, then compute meta-loss.

Original: 10: ¯ θ ← ExpectedUpdate ( θ , d, m ) 11: ˆ θ ← BootstrapTarget (¯ θ , η , x , a , r, x ′ ) 12:
Revised: 10: ¯ θ ← ExpectedUpdate ( θ , d, m ) 11: ˆ θ ← BootstrapTarget (¯ θ , η , x , a , r, x ′ ) 12:

Original: L ( η ) ← ||
Revised: L ( η ) ← ||

Original: [[ˆ θ ]] − ¯ θ ( η ) ||13: η ← Adam ( η , L )
Revised: [[ˆ θ ]] − ¯ θ ( η ) ||13: η ← Adam ( η , L )

Original: 1: input: θ , η , ( x , a , r, x ′ ) 2: for1 , · · · , ℓ do 3: ˜ x , ˜ a ∼ d ( η ) .
Revised: 1: input: θ , η , m 2: for ˜ x ∈ S do 3: for ˜ a ∈

Original:  4: ˜ x ′ , ˜
Revised: A do 4: ˜ x ′ , ˜

Original: r ∼ m (˜ x , ˜ a ) .
Revised: r ∼ m (˜ x , ˜ a ) .

Original: 5: θ ←θ + α [˜ r + γ max  a ′ ∈A ˆ q ( ˜ x ′ ,  a ′ ; θ ) − ˆ q
Revised: 5: ∆ ˜ x , ˜ a ← [˜ r + γ max ˜ a ′ ∈A ˆ q ( ˜ x ′ , ˜ a ′ ; θ ) − ˆ q

Original: (˜ x , ˜ a ; θ )] ∇ θ ˆ q (˜ x , ˜ a ; θ ) 6: ˆ θ ← θ + α [ r + γ max a ′ ∈A ˆ q ( x ′ ,  a ′ ; θ ) − ˆ q ( x, a ; θ )] ∇ θ ˆ q ( x ,  a; θ ) 7: return: ˆ θ
Revised: (˜ x , ˜ a ; θ )] ∇ θ ˆ q (˜ x , ˜ a ; θ ) 6: ¯ θ ← θ + α (cid:80) ˜ x ∈S , ˜ a ∈A π (˜ a | ˜ x ) d ( x ; η )∆ ˜ x , ˜ a 7: return: ¯ θ

Original: 1: input: θ , η , m 2: for ˜ x ∈ S do 3: for ˜ a ∈
Revised: 1: input: θ , η , ( x , a , r, x ′ ) 2: for 1 , · · · , ℓ do 3: ˜ x , ˜ a ∼ d ( η ) .

Original: A do 4: ˜ x ′ , ˜
Revised: 4: ˜ x ′ , ˜

Original: r ∼ m (˜ x , ˜ a ) .
Revised: r ∼ m (˜ x , ˜ a ) [Table] θ

Revised: [Table]

Revised: [˜ r

Revised: [Table]

Original: 5: ∆ ˜ x , ˜ a ← [˜ r + γ max ˜ a ′ ∈A ˆ q ( ˜ x ′ , ˜ a ′ ; θ ) − ˆ q
Revised: γ [Table] ˆ q ( ˜ x [Table]

Revised: ˆ q

Revised: [Table]

Revised: ˜ x

Revised: [Table]

Revised: θ )

Revised: [Table]

Revised: ˆ q

Revised: [Table]

Revised: ˜

Revised: [Table]

Revised: ˜

Revised: [Table]

Revised: ˆ

Revised: [Table]

Revised: θ

Revised: [Table]

Revised: α [ r

Revised: [Table]

Revised: γ max a ′ ∈A ˆ q ( x ′

Revised: [Table]

Revised: a ′ ; θ ) − ˆ q ( x

Revised: [Table]

Original: (˜ x , ˜ a ; θ )] ∇ θ ˆ q (˜ x ,˜ a ; θ ) 6: ¯ θ ← θ + α (cid:80) ˜ x ∈S , ˜ a ∈A π (˜ a | ˜ x ) d ( x ; η )∆ ˜ x , ˜ a 7: return: ¯ θ
Revised: a ; θ )] ∇ θ ˆ q ( x , a ; θ [Table] [Table] return: ˆ θ

Original: A.2 Hyperparameter Selection
Revised: A.2 Hyperparameter Selection

Original: [Table]
Revised: [Table]

Original: To select hyperparameters, we perform a grid search over all possible hyperparameter configurationsfrom Table 1.
Revised: To select hyperparameters [Table] we perform a grid search over all possible hyperparameter configurations from Table 1.

Original: Each configuration is run with 30 random seeds during the selection process.
Revised: Each configuration is run with 30 random seeds during the selection process.

Original: We thenrun the best hyperparameter configurations for an additional 30 seeds and report these as the finalresults.
Revised: We then run the best hyperparameter configurations for an additional 30 seeds and report these as the final results.

Original: [Figure]
Revised: [Figure]

Revised: It is clear that probability is concentrated in states adjacent to the terminal states.

Revised: These states have the greatest change in value when the reward regime switches.

Original: Figure 3 shows snapshots of the probability distribution learned by the API algorithm at variouspoints throughout training.
Revised: Figure 3 shows snapshots of the probability distribution learned by the AP algorithm at various points throughout training.

Original: API begins with a uniform distribution and quickly learns to concentrateprobability in the states nearest to the terminal states.
Revised: AP begins with a uniform distribution and quickly learns to concentrate probability in the states nearest to the terminal states.

Original: As well, it clearly learns the intuition underlyingthe domain-specific distribution described earlier.
Revised: As well, it clearly learns the intuition underlying the domain-specific distribution described earlier.

Original: Very little probability mass is placed in the verticalhallways, where the value of states does not change when the reward regimes changes.
Revised: Very little probability mass is placed in the vertical hallways, where the value of states does not change when the reward regimes changes.

Original: API is clearlylearning that samples from these states is relatively less important than sampling from states whosevalue function does change.
Revised: AP is clearly learning that samples from these states is relatively less important than sampling from states whose value function does change.

Original: A.4 Results for Additional Settings of the Environment
Revised: A.4 Results for Additional Settings of the Environment

Original: Figure 4 shows additional TMaze results for all combinations of ϵ env .
Revised: Figure 4 shows additional TMaze results for all combinations of ϵ env .

Original: The results demonstratethat in all settings, the adaptive algorithm matches the performance of the hand-constructed biaseddistribution while outperforming biased in several settings.
Revised: The results demonstrate that in all settings, the adaptive algorithm matches the performance of the hand-constructed biased distribution while outperforming biased in several settings.

Original: Table TODO INCLUDE REF also reportsthe results of the statistical significance test in each ϵ env and planning steps setting.

Original: A.5 Statistical Significance Tests
Revised: A.5 Statistical Significance Tests

Original: As shown in Figure 2, our metric for evaluating each algorithm was the total reward accumulated overthe course of 100,000 interactions with the TMaze environment.
Revised: As shown in Figure 2, our metric for evaluating each algorithm was the total reward accumulated over the course of 100,000 interactions with the TMaze environment.

Original: To determine whether statisticallysignificant differences exists between the mean results of each algorithm we performed Tukey’smethod.
Revised: To determine whether statistically significant differences exists between the mean results of each algorithm we performed Tukey’s method.

Original: We report the results in the table below.
Revised: We report the results in the table below.

Original: The p-values indicate that in all settings of ϵ env andfor both 3 and 5 steps of planning, API at least matches the performance of the Domain-Specificdistribution.
Revised: The p-values indicate that in all settings of ϵ env and for both 3 and 5 steps of planning, AP at least matches the performance of the Domain-Specific distribution.

Original: In several settings, API exceed the performance of Domain-Specific.
Revised: In several settings, AP exceed the performance of Domain-Specific.

Original: Algorithms Compared
Revised: Algorithms Compared

Original: Difference in Mean Total Reward
Revised: Difference [Table] [Table] [Table]

Original: P-Value
Revised: [Table]

Original: Q-Learning v. Uniform -1375.
Revised: Q-Learning v. [Table] -Specific - [Table]

Revised: [Table]

Original: Q-Learning v. Domain-Specific -1515.
Revised: Q [Table] [Table] [Table]

Revised: Uniform

Original: Q-Learning v. API -1736.
Revised: v. Q -Learning [Table] [Table]

Original: Uniform v. Q-Learning 1375.
Revised: Uniform v. Domain- Specific -

Revised: [Table]

Revised: [Table]

Revised: Uniform

Revised: v. AP

Original: [Table]
Revised: [Table]

Revised: Domain

Revised: [Table]

Revised: Specific v. Q

Revised: [Table]

Revised: Learning

Original:  [Figure] 
Revised: [Table] [Figure] [Table] [Table]

Revised: [Table]

Revised: [Table]

Revised: v. Uniform

Original: [Table] 
Revised: [Table] Q [Table]

Revised: Learning v.

Revised: [Table]

Revised: Specific

Original: [Table]
Revised: [Table]

Revised: Learning v. AP

Original: [Table]
Revised: [Table]

Revised: Uniform v. Q-Learning

Original: [Table]
Revised: [Table]

Revised: Uniform v.

Revised: [Table]

Revised: [Table]

Revised: Uniform v. AP

Original: [Table]
Revised: [Table]

Revised: Domain-Specific v. Q

Revised: [Table]

Revised: Learning

Revised: [Table]

Revised: -Specific v.

Revised: [Table]

Revised: Domain-Specific v.

Revised: [Table]

Revised: [Table]

Revised: Uniform v

Revised: [Table]

Revised: [Table]

Revised: Learning

Revised: [Table]