Original: S KILL T RANSFER VIA P ARTIALLY A MORTIZED H IERARCHICAL P LANNING
Revised: S KILL T RANSFER VIA P ARTIALLY A MORTIZED H IERARCHICAL P LANNING

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: Abstract
Revised: Abstract

Original: To quickly solve new tasks in complex environments, intelligent agents need to build up reusable knowledge.
Revised: To quickly solve new tasks in complex environments, intelligent agents need to build up reusable knowledge.

Original: For example, a learned world model captures knowledge about the environment that applies to new tasks.
Revised: For example, a learned world model captures knowledge about the environment that applies to new tasks.

Original: Similarly, skills capture general behaviors that can apply to new tasks.
Revised: Similarly, skills capture general behaviors that can apply to new tasks.

Original: In this paper, we investigate how these two approaches can be integrated into a single reinforcement learning agent.
Revised: In this paper, we investigate how these two approaches can be integrated into a single reinforcement learning agent.

Original: Speciﬁcally, we leverage the idea of partial amortization for fast adaptation at test time.
Revised: Speciﬁcally, we leverage the idea of partial amortization for fast adaptation at test time.

Original: For this, actions are produced by a policy that is learned over time while the skills it conditions on are chosen using online planning.
Revised: For this, actions are produced by a policy that is learned over time while the skills it conditions on are chosen using online planning.

Original: We demonstrate the beneﬁts of our design decisions across a suite of challenging locomotion tasks and demonstrate improved sample efﬁciency in single tasks as well as in transfer from one task to another, as compared to competitive baselines.
Revised: We demonstrate the beneﬁts of our design decisions across a suite of challenging locomotion tasks and demonstrate improved sample efﬁciency in single tasks as well as in transfer from one task to another, as compared to competitive baselines.

Original: Videos are available at: https://sites.google.com/view/ partial-amortization/home
Revised: Videos are available at: https://sites.google.com/view/ partial-amortization/home

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Humans can effortlessly compose skills, where skills are a sequence of temporally correlated actions and quickly adapt skills learned from one task to another.
Revised: Humans can effortlessly compose skills, where skills are a sequence of temporally correlated actions and quickly adapt skills learned from one task to another.

Original: In order to build re-usable knowledge about the environment, Model-based Reinforcement Learning (MBRL)
Revised: In order to build re-usable knowledge about the environment, Model-based Reinforcement Learning (MBRL)

Original: (Wang et al., 2019) provides an intuitive framework which holds the promise of training agents that generalize to different situations and are sample efﬁcient with respect to number of environment interactions required for training.
Revised: (Wang et al., 2019) provides an intuitive framework which holds the promise of training agents that generalize to different situations and are sample efﬁcient with respect to number of environment interactions required for training.

Original: For temporally composing behaviors, hierarchical reinforcement learning (HRL) (Barto & Mahadevan, 2003) seeks to learn behaviors at different levels of abstraction explicitly
Revised: For temporally composing behaviors, hierarchical reinforcement learning (HRL) (Barto & Mahadevan, 2003) seeks to learn behaviors at different levels of abstraction explicitly

Original: A simple approach for learning the environment dynamics is to learn a world model either directly in the observation space (Chua et al., 2018; Sharma et al., 2019; Wang & Ba, 2019) or in a latent space (Hafner et al., 2019; 2018).
Revised: A simple approach for learning the environment dynamics is to learn a world model either directly in the observation space (Chua et al., 2018; Sharma et al., 2019; Wang & Ba, 2019) or in a latent space (Hafner et al., 2019; 2018).

Original: World models summarize an agent’s experience in the form of learned transition dynamics, and reward models, which are used to learn either parametric policies by amortizing over the entire training experience (Hafner et al., 2019; Janner et al., 2019), or perform online planning as done in Planet (Hafner et al., 2018), and PETS (Chua et al., 2018).
Revised: World models summarize an agent’s experience in the form of learned transition dynamics, and reward models, which are used to learn either parametric policies by amortizing over the entire training experience (Hafner et al., 2019; Janner et al., 2019), or perform online planning as done in Planet (Hafner et al., 2018), and PETS (Chua et al., 2018).

Revised: Amortization here refers to learning a parameterized policy, whose parameters are updated using samples during the training phase, and which can then be directly queried at each state to output an action, during evaluation.

Original: Fully online planning methods such as PETS (Chua et al., 2018) only learn the dynamics (and reward) model and rely on an online search procedure such as Cross-Entropy Method (CEM; Rubinstein, 1997) on the learned models to determine which action to execute next.
Revised: Fully online planning methods such as PETS (Chua et al., 2018) only learn the dynamics (and reward) model and rely on an online search procedure such as Cross-Entropy Method (CEM; Rubinstein, 1997) on the learned models to determine which action to execute next.

Original: Since rollouts from the learned dynamics and reward models are not executed in the actual environment during training, these learned models are sometimes also referred to as imagination models (Hafner et al., 2018; 2019).
Revised: Since rollouts from the learned dynamics and reward models are not executed in the actual environment during training, these learned models are sometimes also referred to as imagination models (Hafner et al., 2018; 2019).

Original: Fully amortized methods such as
Revised: Fully amortized methods such as

Original: Dreamer (Hafner et al., 2019), train a reactive policy with many rollouts from the imagination model.
Revised: Dreamer (Hafner et al., 2019), train a reactive policy with many rollouts from the imagination model.

Original: They then execute the resulting policy in the environment.
Revised: They then execute the resulting policy in the environment.

Original: The beneﬁt of the amortized method is that it becomes better with experience.
Revised: The beneﬁt of the amortized method is that it becomes better with experience.

Original: Amortized policies are also faster.
Revised: Amortized policies are also faster.

Original: An action is computed in one forward pass of the reactive policy as opposed to the potentially expensive search procedure used in CEM.
Revised: An action is computed in one forward pass of the reactive policy as opposed to the potentially expensive search procedure used in CEM.

Original: Additionally, the performance of the amortized method is more consistent as CEM relies on drawing good samples from a random action distribution.
Revised: Additionally, the performance of the amortized method is more consistent as CEM relies on drawing good samples from a random action distribution.

Original: On the other hand, the shortcoming of the amortized policy is generalization.
Revised: On the other hand, the shortcoming of the amortized policy is generalization.

Original: When attempting novel tasks unseen during training, CEM will plan action sequences for the new task, as per the new reward function while a fully amortized method would be stuck with a behaviour optimized for the training tasks.
Revised: When attempting novel tasks unseen during training, CEM will plan action sequences for the new task, as per the new reward function while a fully amortized method would be stuck with a behaviour optimized for the training tasks.

Original: Since it is intractable to perform fully online random shooting based planning in high-dimensional action spaces (Bharadhwaj et al., 2020; Amos & Yarats, 2019), it motivates the question: can we combine online search with amortized policy learning in a meaningful way to learn useful and transferable skills for MBRL?
Revised: Since it is intractable to perform fully online random shooting based planning in high-dimensional action spaces (Bharadhwaj et al., 2020; Amos & Yarats, 2019), it motivates the question: can we combine online search with amortized policy learning in a meaningful way to learn useful and transferable skills for MBRL?

Original: To this end, we propose a partially amortized planning algorithm that temporally composes high-level skills through the Cross-Entropy Method (CEM) (Rubinstein, 1997), and uses these skills to condition a low-level policy that is amortized over the agent’s experience.
Revised: To this end, we propose a partially amortized planning algorithm that temporally composes high-level skills through the Cross-Entropy Method (CEM) (Rubinstein, 1997), and uses these skills to condition a low-level policy that is amortized over the agent’s experience.

Original: Our world model consists of a learned latent dynamics model, and a learned latent reward model.
Revised: Our world model consists of a learned latent dynamics model, and a learned latent reward model.

Original: We have a mutual information (MI) based intrinsic reward objective, in addition to the predicted task rewards that are used to train the low level-policy, while the high level skills are planned through CEM using the learned task rewards.
Revised: We have a mutual information (MI) based intrinsic reward objective, in addition to the predicted task rewards that are used to train the low level-policy, while the high level skills are planned through CEM using the learned task rewards.

Original: We term our approach Learning Skills for Planning ( LSP ).
Revised: We term our approach Learning Skills for Planning ( LSP ).

Original: The key idea of LSP is that the high-level skills are able to abstract out essential information necessary for solving a task, while being agnostic to irrelevant aspects of the environment, such that given a new task in a similar environment, the agent will be able to meaningfully compose the learned skills with very little ﬁne-tuning.
Revised: The key idea of LSP is that the high-level skills are able to abstract out essential information necessary for solving a task, while being agnostic to irrelevant aspects of the environment, such that given a new task in a similar environment, the agent will be able to meaningfully compose the learned skills with very little ﬁne-tuning.

Original: In addition, since the skill-space is low dimensional, we can leverage the beneﬁts of online planning in skill space through CEM, without encountering intractability of using CEM for planning directly in the higher dimensional action space and especially for longer time horizons (Figure 1).
Revised: In addition, since the skill-space is low dimensional, we can leverage the beneﬁts of online planning in skill space through CEM, without encountering intractability of using CEM for planning directly in the higher dimensional action space and especially for longer time horizons (Figure 1).

Original: In summary, our main contributions are developing a partially amortized planning approach for
Revised: In summary, our main contributions are developing a partially amortized planning approach for

Original: MBRL, demonstrating that high-level skills can be temporally composed using this scheme to condition low level policies, and experimentally demonstrating the beneﬁt of LSP over challenging locomotion tasks that require composing different behaviors to solve the task, and beneﬁt in terms of transfer from one quadruped locomotion task to another, with very little adaptation in the target task.
Revised: MBRL, demonstrating that high-level skills can be temporally composed using this scheme to condition low level policies, and experimentally demonstrating the beneﬁt of LSP over challenging locomotion tasks that require composing different behaviors to solve the task, and beneﬁt in terms of transfer from one quadruped locomotion task to another, with very little adaptation in the target task.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: We discuss learning latent dynamics for MBRL, and mutual information skill discovery, that serve as the basic theoretical tools for our approach.
Revised: We discuss learning latent dynamics for MBRL, and mutual information skill discovery, that serve as the basic theoretical tools for our approach.

Original: 2. L EARNING L ATENT D YNAMICS AND B EHAVIORS IN I MAGINATION
Revised: 2. L EARNING L ATENT D YNAMICS AND B EHAVIORS IN I MAGINATION

Original: Latent dynamics models are special cases of world models used in MBRL, that project observations into a latent representation, amenable for planning (Hafner et al., 2019; 2018).
Revised: Latent dynamics models are special cases of world models used in MBRL, that project observations into a latent representation, amenable for planning (Hafner et al., 2019; 2018).

Original: This framework is general as it can model both partially observed environments where sensory inputs can be pixel observations, and fully observable environments, where sensory inputs can be proprioceptive state features.
Revised: This framework is general as it can model both partially observed environments where sensory inputs can be pixel observations, and fully observable environments, where sensory inputs can be proprioceptive state features.

Original: The latent dynamics models we consider in this work, consist of four key components, a representation module p θ ( s t | s t − 1 , a t − 1 , o t ) and an observation module q θ ( o t | s T ) that encode observations and actions to continuous vector-valued latent states s t , a latent forward dynamics module q θ ( s t | s t − 1 , a t − 1 ) that predicts future latent states given only the past states and actions, and a task reward module q θ ( r t | s t ) , that predicts the reward from the environment given the current latent state.
Revised: The latent dynamics models we consider in this work, consist of four key components, a representation module p θ ( s t | s t − 1 , a t − 1 , o t ) and an observation module q θ ( o t | s T ) that encode observations and actions to continuous vector-valued latent states s t , a latent forward dynamics module q θ ( s t | s t − 1 , a t − 1 ) that predicts future latent states given only the past states and actions, and a task reward module q θ ( r t | s t ) , that predicts the reward from the environment given the current latent state.

Original: To learn this model, the agent interacts with the environment and maximizes the following expectation under the dataset of environment interactions D
Revised: To learn this model, the agent interacts with the environment and maximizes the following expectation under the dataset of environment interactions D

Original: = { ( o t , a t , r t ) }
Revised: = { ( o t , a t , r t ) }

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: For optimizing behavior under this latent dynamics model, the agent rolls out trajectories in imagination and estimates the value V ( · ) of the imagined trajectories { s τ , a τ , r τ } t
Revised: For optimizing behavior under this latent dynamics model, the agent rolls out trajectories in imagination and estimates the value V ( · ) of the imagined trajectories { s τ , a τ , r τ } t

Original: + H τ = t through TD ( λ ) estimates as described by Sutton & Barto (2018); Hafner et al.
Revised: + H τ = t through TD ( λ ) estimates as described by Sutton & Barto (2018); Hafner et al.

Original: The agent can either learn a fully amortized policy q φ ( a | s ) as done in Dreamer, by backpropagating through the learned value network v ψ ( · ) or plan online through CEM, for example as in Planet.
Revised: The agent can either learn a fully amortized policy q φ ( a | s ) as done in Dreamer, by backpropagating through the learned value network v ψ ( · ) or plan online through CEM, for example as in Planet.

Original: [Figure]
Revised: [Figure]

Original: 2. M UTUAL I NFORMATION S KILL D ISCOVERY
Revised: 2. M UTUAL I NFORMATION S KILL D ISCOVERY

Original: Some methods for skill discovery have adopted a probabilistic approach that uses the mutual information between skills and future states as an objective (Sharma et al., 2019).
Revised: Some methods for skill discovery have adopted a probabilistic approach that uses the mutual information between skills and future states as an objective (Sharma et al., 2019).

Original: In this approach, skills are represented through a latent variable z upon which a low level policy π ( a | s, z ) is conditioned.
Revised: In this approach, skills are represented through a latent variable z upon which a low level policy π ( a | s, z ) is conditioned.

Original: Given the current state s 0 , skills are sampled from some selection distribution p ( z | s 0 ) .
Revised: Given the current state s 0 , skills are sampled from some selection distribution p ( z | s 0 ) .

Original: The skill conditioned policy is executed under the environment dynamics p d ( s t +1 | s t , a ) resulting in a series of future states abbreviated s (cid:48) :
Revised: The skill conditioned policy is executed under the environment dynamics p d ( s t +1 | s t , a ) resulting in a series of future states abbreviated s (cid:48) :

Original: = { s } .
Revised: = { s } .

Original: Mutual information is deﬁned as:
Revised: Mutual information is deﬁned as:

Original: [Equation]
Revised: [Equation]

Original: It quantiﬁes the reduction in uncertainty about the future states given the skill and vice versa.
Revised: It quantiﬁes the reduction in uncertainty about the future states given the skill and vice versa.

Original: By maximizing the mutual information with respect to the low level policy, the skills are encouraged to produce discernible future states.
Revised: By maximizing the mutual information with respect to the low level policy, the skills are encouraged to produce discernible future states.

Original:  P ARTIAL A MORTIZATION THROUGH H IERARCHY
Revised:  P ARTIAL A MORTIZATION THROUGH H IERARCHY

Original: Our aim is to learn behaviors suitable for solving complex control tasks, and amenable to transfer to different tasks, with minimal ﬁne-tuning.
Revised: Our aim is to learn behaviors suitable for solving complex control tasks, and amenable to transfer to different tasks, with minimal ﬁne-tuning.

Original: To achieve this, we consider the setting of MBRL, where the agent builds up re-usable knowledge of the environment dynamics.
Revised: To achieve this, we consider the setting of MBRL, where the agent builds up re-usable knowledge of the environment dynamics.

Original: For planning, we adopt a partial amortization strategy, such that some aspects of the behavior are re-used over the entire training experience, while other aspects are learned online.
Revised: For planning, we adopt a partial amortization strategy, such that some aspects of the behavior are re-used over the entire training experience, while other aspects are learned online.

Original: We achieve partial amortization by forming high level latent plans and learning a low level policy conditioned on the latent plan.
Revised: We achieve partial amortization by forming high level latent plans and learning a low level policy conditioned on the latent plan.

Original: The three different forms of amortization in planning are described visually through probabilistic graphical models in Figure 2b.
Revised: The three different forms of amortization in planning are described visually through probabilistic graphical models in Figure 2.

Original: We ﬁrst describe the different components of our model, motivate the mutual information based auxiliary objective, and ﬁnally discuss the complete algorithm.
Revised: We ﬁrst describe the different components of our model, motivate the mutual information based auxiliary objective, and ﬁnally discuss the complete algorithm.

Original: World model.
Revised: World model.

Original: Our world model is a latent dynamics model consisting of the components described in section 2.
Revised: Our world model is a latent dynamics model consisting of the components described in section 2.

Original: Low level policy.
Revised: Low level policy.

Original: The low-level policy q φ ( a t | s t , z ) is used to decide which action to execute given the current latent state s t and the currently active skill z .
Revised: The low-level policy q φ ( a t | s t , z ) is used to decide which action to execute given the current latent state s t and the currently active skill z .

Original: Similar to Dreamer, we also train a value model v ψ ( s t ) to estimate the expected rewards the action model achieves from each state s t .
Revised: Similar to Dreamer (Hafner et al., 2019), we also train a value model v ψ ( s t ) to estimate the expected rewards the action model achieves from each state s t .

Original: We estimate value the same way as in equation 6 of Dreamer, balancing bias and variance.
Revised: We estimate value the same way as in equation 6 of Dreamer, balancing bias and variance.

Original: The action model is trained to maximize the estimate of the value, while the value model is trained to ﬁt the estimate of the value that alters as the action model is updated, as done in a typical actor-critic setup (Konda & Tsitsiklis, 2000).
Revised: The action model is trained to maximize the estimate of the value, while the value model is trained to ﬁt the estimate of the value that alters as the action model is updated, as done in a typical actor-critic setup (Konda & Tsitsiklis, 2000).

Original: High level skills.
Revised: High level skills.

Original: In our framework high level skills are continuous random variables that are held for a ﬁxed number K steps.
Revised: In our framework high level skills are continuous random variables that are held for a ﬁxed number K steps.

Original: The high-level skills z are sampled from a skill selection distribution p ( z 1: (cid:100) H/K (cid:101) | ζ ) =
Revised: The high-level skills z are sampled from a skill selection distribution p ( z 1: (cid:100) H/K (cid:101) | ζ ) =

Original: N ( µ, Σ) which is optimized for task performance through CEM.
Revised: N ( µ, Σ) which is optimized for task performance through CEM.

Original: Here, H denotes the planning horizon.
Revised: Here, H denotes the planning horizon.

Original: For the sake of notational convenience we denote z 1: (cid:100) H/K (cid:101) as z .
Revised: For the sake of notational convenience we denote z 1: (cid:100) H/K (cid:101) as z .

Original: Let ( j ) denote the j th CEM iteration.
Revised: Let ( j ) denote the j th CEM iteration.

Original: We ﬁrst sample G skills { z ( g ) } Gg =1 ∼ p ( z | ζ ( j ) ) , execute G parallel imaginary rollouts of horizon H in the learned model with the skill-conditioned policy q φ ( a t | s t , z ( g ) ) .
Revised: We ﬁrst sample G skills { z ( g ) } Gg =1 ∼ p ( z | ζ ( j ) ) , execute G parallel imaginary rollouts of horizon H in the learned model with the skill-conditioned policy q φ ( a t | s t , z ( g ) ) .

Original: Instead of evaluating rollouts based only on the sum of rewards, we utilize the value network and compute value estimates { V g } Gg =1 .
Revised: Instead of evaluating rollouts based only on the sum of rewards, we utilize the value network and compute value estimates { V g } Gg =1 .

Original: We sort { V g } Gg =1 , choose the top M values, and use the corresponding skills to update the sampling distribution parameters as ζ ( j +1) =
Revised: We sort { V g } Gg =1 , choose the top M values, and use the corresponding skills to update the sampling distribution parameters as ζ ( j +1) =

Original: ( µ ( j +1) , Σ ( j +1) )
Revised: ( µ ( j +1) , Σ ( j +1) )

Original: [Equation]
Revised: [Equation]

Original: 3. O VERALL A LGORITHM
Revised: 3. O VERALL A LGORITHM

Original: Our overall algorithm consists of the three phases typical in a MBRL pipeline, that are performed iteratively.
Revised: Our overall algorithm consists of the three phases typical in a MBRL pipeline, that are performed iteratively.

Original: The complete algorithm is shown in Algorithm 1.
Revised: The complete algorithm is shown in Algorithm 1.

Original: The sub-routine for CEM planning that gets called in Algorithm 1 is described in Algorithm 2.
Revised: The sub-routine for CEM planning that gets called in Algorithm 1 is described in Algorithm 2.

Original: Model Learning.
Revised: Model Learning.

Original: We sample a batch of tuples from the dataset of environment interactions { ( a t , o t , r t ) } k + L t = k ∼ D , compute the latent states s t ∼ p
Revised: We sample a batch of tuples from the dataset of environment interactions { ( a t , o t , r t ) } k + L t = k ∼ D , compute the latent states s t ∼ p

Original: θ ( s t | s t − 1 , a t − 1 , o t ) , and use the resulting data to update the models p θ ( s t | s t - 1 , a t - 1 , o t ) , q θ ( s t | s t - 1 , a t - 1 ) , and q θ ( r t | s t ) through the variational information bottleneck (VIB)
Revised: θ ( s t | s t − 1 , a t − 1 , o t ) , and use the resulting data to update the models p θ ( s t | s t - 1 , a t - 1 , o t ) , q θ ( s t | s t - 1 , a t - 1 ) , and q θ ( r t | s t ) through the variational information bottleneck (VIB)

Original: (Tishby et al., 2000; Alemi et al., 2016) objective as in equation 13 of Hafner et al.
Revised: (Tishby et al., 2000; Alemi et al., 2016) objective as in equation 13 of Hafner et al.

Original: and as described in section 2.
Revised: and as described in section 2.

Original: Behavior Learning.
Revised: Behavior Learning.

Original: Here, we roll out the low-level policy q φ ( a t | s t , z ) in the world model and use the state transitions and predicted rewards to optimize the parameters of the policy φ , the skill distribution ζ , the value model ψ , and the backward skill predictor χ . 
Revised: Here, we roll out the low-level policy q φ ( a t | s t , z ) in the world model and use the state transitions and predicted rewards to optimize the parameters of the policy φ , the skill distribution ζ , the value model ψ , and the backward skill predictor χ . The backward skill predictor predicts the skill z given latent rollouts { s } .

Original: Environment Interaction.
Revised: Environment Interaction.

Original: This step is to collect data in the actual environment for updating the model parameters.
Revised: This step is to collect data in the actual environment for updating the model parameters.

Original: Using Model-Predictive Control (MPC), we re-sample high-level skills from the optimized p ζ ( z ) every K steps, and execute the low-level policy q φ ( a t | s t , z ) , conditioned on the currently active skill z .
Revised: Using Model-Predictive Control (MPC), we re-sample high-level skills from the optimized p ζ ( z ) every K steps, and execute the low-level policy q φ ( a t | s t , z ) , conditioned on the currently active skill z .

Original: Hence, the latent plan has a lower temporal resolution as compared to the low level policy.
Revised: Hence, the latent plan has a lower temporal resolution as compared to the low level policy.

Original: This helps us perform temporally abstracted exploration easily in the skill space.
Revised: This helps us perform temporally abstracted exploration easily in the skill space.

Original: We store the (observation, action, reward) tuples in the dataset D .
Revised: We store the (observation, action, reward) tuples in the dataset D .

Original: 3. M UTUAL I NFORMATION S KILL O BJECTIVE
Revised: 3. M UTUAL I NFORMATION S KILL O BJECTIVE

Original: Merely conditioning the low level policy q φ ( a t | s t , z ) on the skill z is not sufﬁcient as it is prone to ignoring it.
Revised: Merely conditioning the low level policy q φ ( a t | s t , z ) on the skill z is not sufﬁcient as it is prone to ignoring it.

Original: Hence, we incorporate maximization of the mutual information (MI) between the latent skills z and the sequence of states { s } as an auxiliary objective.
Revised: Hence, we incorporate maximization of the mutual information (MI) between the latent skills z and the sequence of states { s } as an auxiliary objective.

Original: In this paper, we make use of imagination rollouts to estimate the mutual information under the agent’s learned dynamics model.
Revised: In this paper, we make use of imagination rollouts to estimate the mutual information under the agent’s learned dynamics model.

Original: We decompose the mutual information in terms of skill uncertainty reduction MI ( z, { s }| s 0 ) =
Revised: We decompose the mutual information in terms of skill uncertainty reduction MI ( z, { s }| s 0 ) =

Original: H ( z | s 0 ) − H ( z |{ s } , s 0 ) .
Revised: H ( z | s 0 ) − H ( z |{ s } , s 0 ) .

Original: Estimating MI ( z, { s }| s 0 ) .
Revised: Estimating MI ( z, { s }| s 0 ) .

Original: Explicitly writing out the entropy terms, we have
Revised: Explicitly writing out the entropy terms, we have

Original: [Equation]
Revised: [Equation]

Original: In this case we need a tractable approximation to the skill posterior p ( z | s 0 , s (cid:48) ) .
Revised: In this case we need a tractable approximation to the skill posterior p ( z | s 0 , s (cid:48) ) .

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Here the latter term is a KL divergence and must hence be positive, providing a lower bound for MI .
Revised: Here the latter term is a KL divergence and must hence be positive, providing a lower bound for MI .

Original: [Equation]
Revised: [Equation]

Original: We parameterize q ( z |{ s } , s 0 ) with χ , i.e. q χ ( z |{ s } , s 0 ) , and call it the backward skill predictor, as it predicts the skill
Revised: We parameterize q ( z |{ s } , s 0 ) with χ , i.e. q χ ( z |{ s } , s 0 ) , and call it the backward skill predictor, as it predicts the skill

Original: z given latent rollouts { s } .
Revised: z given latent rollouts { s } .

Original: It is trained through standard supervised learning to maximize the likelihood of imagined rollouts E p ( z, { s } ,s 0 )
Revised: It is trained through standard supervised learning to maximize the likelihood of imagined rollouts E p ( z, { s } ,s 0 )

Original: [log q ( z |{ s } , s 0 )] .
Revised: [log q ( z |{ s } , s 0 )] .

Original: This mutual information objective is only a function of the policy through the ﬁrst term and hence we use it as the intrinsic reward for the agent r i
Revised: This mutual information objective is only a function of the policy through the ﬁrst term and hence we use it as the intrinsic reward for the agent r i

Original: = log q
Revised: = log q

Original: ( z |{ s } , s 0 ) .
Revised: ( z |{ s } , s 0 ) .

Original: The second term E s 

Revised: The second term E s

Original: [ H [ p ( z | s 0 )]] is the entropy of the skill selection distribution.
Revised: [ H [ p ( z | s 0 )]] is the entropy of the skill selection distribution.

Original: When skills begin to specialize, the CEM distribution will naturally decrease in entropy and so we add Gaussian noise (cid:15) to the mean of the CEM-based skill distribution, µ ← µ + (cid:15) where (cid:15) ∼ N (0 , I σ ) .
Revised: When skills begin to specialize, the CEM distribution will naturally decrease in entropy and so we add Gaussian noise (cid:15) to the mean of the CEM-based skill distribution, µ ← µ + (cid:15) where (cid:15) ∼ N (0 , I σ ) .

Original: By doing this we lower bound the entropy of the skill selection distribution.
Revised: By doing this we lower bound the entropy of the skill selection distribution.

Original: Algorithm
Revised: Algorithm

Revised: Learning Skills for Planning

Original: Initialize dataset D with S random seed episodes.
Revised: Initialize dataset D with S random seed episodes.

Original: Initialize neural network parameters θ, φ, ψ
Revised: Initialize neural network parameters θ, φ, ψ

Original: Deﬁne skill duration K , plan horizon H , CEM population size G , MaxCEMiter, skill noise (cid:15) while not converged do for update step c = 1 ..
Revised: Deﬁne skill duration K , plan horizon H , CEM population size G , MaxCEMiter, skill noise (cid:15) while not converged do for update step c = 1 ..

Original: C do 
Revised: C do //

Original: Model learning
Revised: Model learning

Original: Draw B data sequences { ( a t , o t , r t ) } k + L t = k ∼ D .
Revised: Draw B data sequences { ( a t , o t , r t ) } k + L t = k ∼ D .

Original: Compute model states s t ∼ p θ ( s t | s t − 1 , a t − 1 , o t ) .
Revised: Compute model states s t ∼ p θ ( s t | s t − 1 , a t − 1 , o t ) .

Original: S ← S ∪ { s t }
Revised: S ← S ∪ { s t }

Original: Update θ using representation learning.
Revised: Update θ using representation learning.

Original:  Behavior learning
Revised: / Behavior learning

Original: ζ = ( µ, Σ) ← CEM ( S , MaxCEMiter, G , H , K , ζ (0) ); Add noise µ = µ + (cid:15)
Revised: ζ = ( µ, Σ) ← CEM ( S , MaxCEMiter, G , H , K , ζ (0) ); Add noise µ = µ + (cid:15)

Original: Compute rewards R = E (cid:0) q θ ( r τ | s τ ) (cid:1) with the optimized CEM distribution ζ
Revised: Compute rewards R = E (cid:0) q θ ( r τ | s τ ) (cid:1) with the optimized CEM distribution ζ

Original: Compute the corresponding intrinsic rewards r
Revised: Compute the corresponding intrinsic rewards r

Original: i
Revised: i

Original: Store the corresponding states into K − sized sequences {{ s } k } (cid:100) H/K (cid:101) k
Revised: Store the corresponding states into K − sized sequences {{ s } k } (cid:100) H/K (cid:101) k

Original: = (for χ )
Revised: = (for χ )

Original: Use total rewards R + r i to form value estimates, and update φ, ψ, χ //
Revised: Use total rewards R + r i to form value estimates, and update φ, ψ, χ //

Original: Environment 
Revised: Environment interaction

Original: interaction o 1 ← env.reset() for time step t = 

Revised: o 1 ← env.reset() for time step t =

Original: T − 1 do //
Revised: T − 1 do //

Original: MPC in z space.
Revised: MPC in z space.

Original: Resample skill every K timesteps
Revised: Resample skill every K timesteps

Original: if t % K == 

Revised: if t % K ==

Original: then
Revised: then

Original: Sample skill /K (cid:100) /K
Revised: Sample skill z 1: (cid:100) H/K (cid:101) ∼ p ( z 1: (cid:100) H/K (cid:101) | ζ )

Original: Choose the ﬁrst skill,
Revised: Choose the ﬁrst skill z

Original: =
Revised: =

Revised: z

Original: Compute s  ∼ p
Revised: Compute s t ∼ p

Original:  ( s t
Revised: θ ( s t

Original:  s t − 1 , a t − 1 , o t ) from history, choose a t 
Revised: | s t − 1 , a t − 1 , o t ) from history, choose a t ∼ q φ

Original: q φ ( a t 
Revised: ( a t |

Original: s t , z ) r t , o t +1 ← env.step Add
Revised: s t , z ) r t , o t +1 ← env.step( a t )

Original:  experience to dataset D 
Revised: Add experience to dataset D ← D ∪ { ( o t , a t , r t )

Original: D o , a Tt 
Revised: Tt =1 } .

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We perform experimental evaluation over locomotion tasks based on the DeepMind Control Suite framework (Tassa et al., 2018) to understand the following questions:
Revised: We perform experimental evaluation over locomotion tasks based on the DeepMind Control Suite framework (Tassa et al., 2018) to understand the following questions:

Original: • Does LSP learn useful skills and compose them appropriately to succeed in individual tasks?
Revised: • Does LSP learn useful skills and compose them appropriately to succeed in individual tasks?

Original: • Does LSP adapt to a target task with different environment reward functions quickly, after being pre-trained on another task?
Revised: • Does LSP adapt to a target task with different environment reward functions quickly, after being pre-trained on another task?

Original: To answer these, we perform experiments on locomotion tasks, using agents with different dynamics - Quadruped, Walker, Cheetah, and Hopper, and environments where either pixel observations or proprioceptive features are available to the agent.
Revised: To answer these, we perform experiments on locomotion tasks, using agents with different dynamics - Quadruped, Walker, Cheetah, and Hopper, and environments where either pixel observations or proprioceptive features are available to the agent.

Original: Our experiments consist of evaluation in single tasks, in transfer from one task to another, ablation studies, and visualization of the learned skills.
Revised: Our experiments consist of evaluation in single tasks, in transfer from one task to another, ablation studies, and visualization of the learned skills.

Original: 4. S ETUP
Revised: 4. S ETUP

Original: Baselines.
Revised: Baselines.

Original: We consider Dreamer (Hafner et al., 2019), which is a state of the art model-based RL algorithm with fully amortized policy learning, as the primary baseline, based on its opensource tensorﬂow2 implementation.
Revised: We consider Dreamer (Hafner et al., 2019), which is a state of the art model-based RL algorithm with fully amortized policy learning, as the primary baseline, based on its open-source tensorﬂow2 implementation.

Revised: We consider a Hierarchical RL (HRL) baseline, HIRO (Nachum et al., 2018) that trains a high level amortized policy (as opposed to high level planning).

Revised: For consistency, we use the same intrinsic reward for HIRO as our method.

Original: We consider another baseline, named Random Skills that has a hierarchical structure exactly as LSP but the skills are sampled randomly and there is no planning at the level of skills.
Revised: We consider two other baselines, named Random Skills that has a hierarchical structure exactly as LSP but the skills are sampled randomly and there is no planning at the level of skills, and RandSkillsInit that is similar to Random Skills but does not include the intrinsic rewards (this is essentially equivalent to Dreamer with an additional random skill input)

Original: This is to help understand the utility of the learned skills.
Revised: These two baselines are to help understand the utility of the learned skills.

Revised: For the transfer experiments, we consider an additional baseline, a variant of our method that keeps the low-level policy ﬁxed in the transfer environment.

Original: All results are over three random seeds.
Revised: All results are over three random seeds.

Original: Environments.
Revised: Environments.

Original: We consider challenging locomotion environments from DeepMind Control
Revised: We consider challenging locomotion environments from DeepMind Control

Original: Suite (Tassa et al., 2018) for evaluation, that require learning walking, running, and hopping gaits which can be achieved by temporally composing skills.
Revised: Suite (Tassa et al., 2018) for evaluation, that require learning walking, running, and hopping gaits which can be achieved by temporally composing skills.

Original: In the Quadruped GetUp Walk task, aa quadruped must learn to stand up from a randomly initialized position that is sometimes upside down, and walk on a plane without toppling, while in Quadruped Reach , the quadruped agent must walk in order to reach a particular goal location.
Revised: In the Quadruped GetUp Walk task, a quadruped must learn to stand up from a randomly initialized position that is sometimes upside down, and walk on a plane without toppling, while in Quadruped Reach , the quadruped agent must walk in order to reach a particular goal location.

Original: In Quadruped Run , the same quadruped agent must run as fast as possible, with higher rewards for faster speed.
Revised: In Quadruped Run , the same quadruped agent must run as fast as possible, with higher rewards for faster speed.

Original: In Cheetah Run , and Walker Run , the cheetah and walker agents must run as fast as possible.
Revised: In Cheetah Run , and Walker Run , the cheetah and walker agents must run as fast as possible.

Original: In Hopper Hop , a one legged hopper must hop in the environment without toppling.
Revised: In Hopper Hop , a one legged hopper must hop in the environment without toppling.

Original: It is extremely challenging to maitain stability of this agent.
Revised: It is extremely challenging to maintain stability of this agent.

Original: 4. T OY T EST
Revised: 4. T OY T EST

Original: We ﬁrst validate our approach on a simple toy environment consisting of a 7D point mass with drag forces trying to navigate to high reward regions of the state space.
Revised: We ﬁrst validate our approach on a simple toy environment consisting of a 7D point mass with drag forces trying to navigate to high reward regions of the state space.

Original: The reward is characterized by distance to a randomly sampled goal which takes small random movements throughout the episode.
Revised: The reward is characterized by distance to a randomly sampled goal which takes small random movements throughout the episode.

Original: As shown in Figure 1, high level planning of LSP rapidly adapts to generate a diverse set of high reward temporally correlated trajectories as compared to a random shooting (CEM) planner for actions.
Revised: As shown in Figure 1, high level planning of LSP rapidly adapts to generate a diverse set of high reward temporally correlated trajectories as compared to a random shooting (CEM) planner for actions.

Original: This demonstrates the intractability of online planning in high-dimensional action spaces, and in part motivates our approach of using it to plan at the level of low-dimensional high level skills.
Revised: This demonstrates the intractability of online planning in high-dimensional action spaces, and in part motivates our approach of using it to plan at the level of low-dimensional high level skills.

Original: 4. S OLVING SINGLE LOCOMOTION TASKS .
Revised: 4. S OLVING SINGLE LOCOMOTION TASKS .

Original: In Figure 3 we evaluate our approach LSP in comparison to the fully amortized baseline Dreamer , and the Random Skills baseline on a suite of challenging locomotion tasks.
Revised: In Figure 3 we evaluate our approach LSP in comparison to the fully amortized baseline Dreamer , and the Random Skills baseline on a suite of challenging locomotion tasks.

Original: Although the environments have a single task objective, in order to achieve high rewards, the agents need to learn different walking gaits (Quadruped Walk, Walker Walk), running gaits (Quadruped Run, Walker Run), and hopping gaits (Hopper Hop) and compose learned skills appropriately for locomotion.
Revised: Although the environments have a single task objective, in order to achieve high rewards, the agents need to learn different walking gaits (Quadruped Walk, Walker Walk), running gaits (Quadruped Run, Walker Run), and hopping gaits (Hopper Hop) and compose learned skills appropriately for locomotion.

Original: From the results, it is evident that LSP either outperforms Dreamer or is competitive to it on all the environments.
Revised: From the results, it is evident that LSP either outperforms Dreamer or is competitive to it on all the environments.

Original: In addition, we observe that LSP signiﬁcantly outperforms the Random Skills baseline, indicating that learning skills and planning over them is important to succeed in these locomotion tasks.
Revised: In addition, we observe that LSP signiﬁcantly outperforms the Random Skills baseline, indicating that learning skills and planning over them is important to succeed in these locomotion tasks.

Original: 4. T RANSFER FROM ONE TASK TO ANOTHER .
Revised: 4. T RANSFER FROM ONE TASK TO ANOTHER .

Original: In Figure 4 we show results for a quadruped agent that is pretrained on one task and must transfer to a different task with similar environment dynamics but different task description.
Revised: In Figure 4 we show results for a quadruped agent that is pretrained on one task and must transfer to a different task with similar environment dynamics but different task description.

Original: Quadruped GetUp Walk → Reach Goal.
Revised: Quadruped GetUp Walk → Reach Goal.

Original: The quadruped agent is pre-trained on the task of  walking as fast as possible while remaining upright.
Revised: The quadruped agent is pre-trained on the task of standing up from a randomly initialized position that is sometimes upside down,and walking on a plane without toppling.

Original: The transfer task consists of walking to reach a goal, and environment rewards are speciﬁed in terms of distance to goal.
Revised: The transfer task consists of walking to reach a goal, and environment rewards are speciﬁed in terms of distance to goal.

Original: The agent is randomly initialized and is sometimes initialized upside down, such that it must learn to get upright and then start walking towards the goal.
Revised: The agent is randomly initialized and is sometimes initialized upside down, such that it must learn to get upright and then start walking towards the goal.

Original: We see that LSP can adapt much quickly to the transfer task, achieving a reward of 500 only after 70 , 000 steps, while Dreamer requires 130 , 000 steps to achieve the same reward, indicating sample efﬁcient transfer of learned skills.
Revised: We see that LSP can adapt much quickly to the transfer task, achieving a reward of 500 only after 70 , 000 steps, while Dreamer requires 130 , 000 steps to achieve the same reward, indicating sample efﬁcient transfer of learned skills.

Original: Quadruped Walk → Reach Goal.
Revised: Quadruped Walk → Reach Goal.

Original: The quadruped agent is randomly initialized, but it is ensured that it is upright at initialization.
Revised: The quadruped agent is randomly initialized, but it is ensured that it is upright at initialization.

Original: In this setting, after pre-training, we re-label the value of rewards in the replay buffer of both the baseline Dreamer , and LSP with the reward function of the target
Revised: In this setting, after pre-training, we re-label the value of rewards in the replay buffer of both the baseline Dreamer , and LSP with the reward function of the target

Original: Quadruped Reach Goal environment.
Revised: Quadruped Reach Goal environment.

Original: To do this, we consider each tuple in the replay buffer of imagined trajectories during pre-training, and change the reward labels to the reward value obtained by querying the reward function of the target task at the corresponding state and action of the tuple.
Revised: To do this, we consider each tuple in the replay buffer of imagined trajectories during pre-training, and change the reward labels to the reward value obtained by querying the reward function of the target task at the corresponding state and action of the tuple.

Original: From the plot in Figure 4, we see that LSP is able to quickly bootstrap learning from the re-labeled
Revised: From the plot in Figure 4, we see that LSP is able to quickly bootstrap learning from the re-labeled

Original: Replay Buffer and achieve better target adaptation than the baseline.
Revised: Replay Buffer and achieve better target adaptation than the baseline.

Original: 4. M UTUAL I NFORMATION A BLATION STUDY
Revised: 4. M UTUAL I NFORMATION A BLATION STUDY

Original: In order to better understand the beneﬁt of the mutual information skill objective, we compare performance against a baseline that is equivalent to LSP but does not use the intrinsic reward to train the low level policy.
Revised: In order to better understand the beneﬁt of the mutual information skill objective, we compare performance against a baseline that is equivalent to LSP but does not use the intrinsic reward to train the low level policy.

Original: We call this ablation baseline that does not have mutual information maximization between skills and states, as no_MI .
Revised: We call this ablation baseline that does not have mutual information maximization between skills and states, as no_MI .

Original: We show the respective reward curves for the
Revised: We show the respective reward curves for the

Original: Quadruped GetUp Walk task in Figure 5.
Revised: Quadruped GetUp Walk task in Figure 5.

Original: Without the mutual information objective, LSP learns less quickly but still faster than the baseline Dreamer .
Revised: Without the mutual information objective, LSP learns less quickly but still faster than the baseline Dreamer .

Original: This emphasizes the necessity of the MI skill objective in section 3.2 and suggests that merely conditioning the low-level policy on the learned skills is still effective to some extent but potentially suffers from the low-level policy learning to ignore them.
Revised: This emphasizes the necessity of the MI skill objective in section 3.2 and suggests that merely conditioning the low-level policy on the learned skills is still effective to some extent but potentially suffers from the low-level policy learning to ignore them.

Original: 4. V ISUALIZATION OF THE LEARNED SKILLS
Revised: 4. V ISUALIZATION OF THE LEARNED SKILLS

Original: In Figure 6, we visualize learned skills of LSP , while transferring from the Quadruped Walk to the Quadruped Reach task.
Revised: In Figure 6, we visualize learned skills of LSP , while transferring from the Quadruped Walk to the Quadruped Reach task.

Original: Each sub-ﬁgure (with composited images) corresponds to a different trajectory rolled out from the same initial state.
Revised: Each sub-ﬁgure (with composited images) corresponds to a different trajectory rolled out from the same initial state.

Original: It is evident that the learned skills are reasonably diverse and useful in the transfer task.
Revised: It is evident that the learned skills are reasonably diverse and useful in the transfer task.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Skill discovery.
Revised: Skill discovery.

Original: Some RL algorithms explicitly try to learn task decompositions in the form of re-usable skills, which are generally formulated as temporally abstracted actions (Sutton et al., 1999).
Revised: Some RL algorithms explicitly try to learn task decompositions in the form of re-usable skills, which are generally formulated as temporally abstracted actions (Sutton et al., 1999).

Original: Most recent skill discovery algorithms seek to maximize the mutual information between skills and input observations (Gregor et al., 2016; Florensa et al., 2017), sometimes resulting in an unsupervised diversity maximization objective (Eysenbach et al., 2018; Sharma et al., 2019).
Revised: Most recent skill discovery algorithms seek to maximize the mutual information between skills and input observations (Gregor et al., 2016; Florensa et al., 2017), sometimes resulting in an unsupervised diversity maximization objective (Eysenbach et al., 2018; Sharma et al., 2019).

Original: DADS (Sharma et al., 2019) is an unsupervised skill discovery algorithm for learning diverse skills with a skill-transition dynamics model, but does not learn a world model for low-level actions and observations, and hence cannot learn through imagined rollouts and instead requiring many environment rollouts with different sampled skills.
Revised: DADS (Sharma et al., 2019) is an unsupervised skill discovery algorithm for learning diverse skills with a skill-transition dynamics model, but does not learn a world model for low-level actions and observations, and hence cannot learn through imagined rollouts and instead requiring many environment rollouts with different sampled skills.

Original: Transfer in RL.
Revised: Hierarchical RL.

Revised: Hierarchical RL (HRL) (Barto & Mahadevan, 2003) methods decompose a complex task to sub-tasks and solve each task by optimizing a certain objective function.

Revised: HIRO (Nachum et al., 2018) learns a high-level policy and a low-level policy and computes intrinsic rewards for training the low-level policy through sub-goals speciﬁed as part of the state-representation the agent observes.

Revised: Some other algorithms follow the options framework (Sutton et al., 1999; Bacon et al., 2017), where options correspond to temporal abstractions that need specifying some termination conditions.

Revised: In practice, it is difﬁcult to learn meaningful termination conditions without additional regularization (Harb et al., 2017).

Revised: These HRL approaches are inherently speciﬁc to the tasks being trained on, and do not necessarily transfer to new domains, even with similar dynamics.

Revised: Transfer in RL.

Original: Multiple previous works have investigated the problem of transferring policies to different environments.
Revised: Multiple previous works have investigated the problem of transferring policies to different environments.

Original: Progressive Networks (Rusu et al., 2016) bootstrap knowledge from previously learned tasks by avoiding catastrophic forgetting of the learned models, (Byravan et al., 2020) perform model-based value estimation for learning an amortized policy and transfer to tasks with different speciﬁed reward functions keeping the dynamics the same, while Plan2Explore (Sekar et al., 2020)
Revised: Progressive Networks (Rusu et al., 2016) bootstrap knowledge from previously learned tasks by avoiding catastrophic forgetting of the learned models, (Byravan et al., 2020) perform model-based value estimation for learning an amortized policy and transfer to tasks with different speciﬁed reward functions keeping the dynamics the same, while Plan2Explore (Sekar et al., 2020)

Original: ﬁrst learns a global world model without task rewards through a self-supervised objective, and given a user-speciﬁed reward function at test time, quickly adapts to it.
Revised: ﬁrst learns a global world model without task rewards through a self-supervised objective, and given a user-speciﬁed reward function at test time, quickly adapts to it.

Original: In contrast to these, several meta RL approaches learn policy parameters that generalize well with little ﬁne-tuning (often in the form of gradient updates) to target environments (Finn et al., 2017; Xu et al., 2018;
Revised: In contrast to these, several meta RL approaches learn policy parameters that generalize well with little ﬁne-tuning (often in the form of gradient updates) to target environments (Finn et al., 2017; Xu et al., 2018;

Original: Wang et al., 2016; Rakelly et al., 2019; Yu et al., 2019).
Revised: Wang et al., 2016; Rakelly et al., 2019; Yu et al., 2019).

Original: Amortization for planning.
Revised: Amortization for planning.

Original: Most current MBRL approaches use some version of the ‘CrossEntropy Method’ (CEM) for doing a random population based search of plans given the current model (Wang & Ba, 2019; Hafner et al., 2018).
Revised: Most current MBRL approaches use some version of the ‘CrossEntropy Method’ (CEM) for doing a random population based search of plans given the current model (Wang & Ba, 2019; Hafner et al., 2018).

Original: Some papers (Sharma et al., 2019) use other population based search approaches like Model-Predictive Path Integral (MPPI)
Revised: Some papers (Sharma et al., 2019) use other population based search approaches like Model-Predictive Path Integral (MPPI)

Original: (Williams et al., 2016) for planning.
Revised: (Williams et al., 2016) for planning.

Original: A recent paper (Okada & Taniguchi, 2019) discusses how in the control as inference framework, the two approaches, CEM, and MPPI are very similar, and differ only with respect to the scoring function r ( τ ) .
Revised: A recent paper (Okada & Taniguchi, 2019) discusses how in the control as inference framework, the two approaches, CEM, and MPPI are very similar, and differ only with respect to the scoring function r ( τ ) .

Original: These online non-amortized planning approaches are typically very expensive in high-dimensional action spaces.
Revised: These online non-amortized planning approaches are typically very expensive in high-dimensional action spaces.

Original: Although (Wang & Ba, 2019) introduces the idea of performing the CEM search in the parameter space of a distilled policy, it still is very costly and requires a lot of samples for convergence.
Revised: Although (Wang & Ba, 2019) introduces the idea of performing the CEM search in the parameter space of a distilled policy, it still is very costly and requires a lot of samples for convergence.

Original: To mitigate these issues, some recent approaches have combined gradient-descent based planning with CEM (Bharadhwaj et al., 2020; Amos & Yarats,2019).
Revised: To mitigate these issues, some recent approaches have combined gradient-descent based planning with CEM (Bharadhwaj et al., 2020; Amos & Yarats,

Original: In contrast, (Janner et al., 2019; Hafner et al., 2019) fully amortize learned policies over the entire training experience, which is fast even for high-dimensional action spaces, but cannot directly transfer to new environments with different dynamics and reward functions.
Revised: In contrast, (Janner et al., 2019; Hafner et al., 2019) fully amortize learned policies over the entire training experience, which is fast even for high-dimensional action spaces, but cannot directly transfer to new environments with different dynamics and reward functions.

Original: We combined the best of both approaches by using CEM to plan online for high-level skills (of low dimensionality) and amortize the skill conditioned policy for low-level actions (of higher dimensionality).
Revised: We combined the best of both approaches by using CEM to plan online for high-level skills (of low dimensionality) and amortize the skill conditioned policy for low-level actions (of higher dimensionality).

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: In this paper, we analyzed the implications of partial amortization with respect to sample efﬁciency and overall performance on a suite of locomotion and transfer tasks.
Revised: In this paper, we analyzed the implications of partial amortization with respect to sample efﬁciency and overall performance on a suite of locomotion and transfer tasks.

Original: We speciﬁcally focused on the setting where partial amortization is enforced through a hierarchical planning model consisting of a fully amortized low-level policy and a fully online high level skill planner.
Revised: We speciﬁcally focused on the setting where partial amortization is enforced through a hierarchical planning model consisting of a fully amortized low-level policy and a fully online high level skill planner.

Original: Through experiments in both state-based and image-based environments we demonstrated the efﬁcacy of our approach in terms of planning for useful skills and executing high-reward achieving policies conditioned on those skills, as evaluated by sample efﬁciency (measured by number of environment interactions) and asymptotic performance (measured by cumulative rewards and success rate).
Revised: Through experiments in both state-based and image-based environments we demonstrated the efﬁcacy of our approach in terms of planning for useful skills and executing high-reward achieving policies conditioned on those skills, as evaluated by sample efﬁciency (measured by number of environment interactions) and asymptotic performance (measured by cumulative rewards and success rate).

Original: One key limitation of our algorithm is that CEM planning is prohibitive in high-dimensional action spaces, and so we cannot have a very high dimensional skill-space for planning with CEM, that might be necessary for learning more expressive/complex skills in real-world robot control tasks.
Revised: One key limitation of our algorithm is that CEM planning is prohibitive in high-dimensional action spaces, and so we cannot have a very high dimensional skill-space for planning with CEM, that might be necessary for learning more expressive/complex skills in real-world robot control tasks.

Original: One potential direction of future work is to incorporate amortized learning of skill policies during training, and use CEM for online planning of skills only during inference.
Revised: One potential direction of future work is to incorporate amortized learning of skill policies during training, and use CEM for online planning of skills only during inference.

Original: Another direction could be to incorporate gradient-descent based planning instead of a random search procedure as CEM, but avoiding local optima in skill planning would be a potential 
Revised: Another direction could be to incorporate gradient-descent based planning instead of a random search procedure as CEM, but avoiding local optima in skill planning would be a potential challenge for gradient descent planning.

Original: A A PPENDIX A. S ETTINGS AND H YPERPARAMETERS
Revised: A A PPENDIX A. S ETTINGS AND H YPERPARAMETERS

Original: Our method is based on the tensorﬂow2 implementation of Dreamer (Hafner et al., 2019) and retains most of the original hyperparameters for policy and model learning.
Revised: Our method is based on the tensorﬂow2 implementation of Dreamer (Hafner et al., 2019) and retains most of the original hyperparameters for policy and model learning.

Original: Training is performed more frequently, in particular 1 training update of all modules is done every 5 environment steps.
Revised: Training is performed more frequently, in particular 1 training update of all modules is done every 5 environment steps.

Original: This marginally improves the training performance and is used in the dreamer baseline as well as our method.
Revised: This marginally improves the training performance and is used in the dreamer baseline as well as our method.

Original: For feature-based experiments we replace the convolutional encoder and decoder with 2-layer
Revised: For feature-based experiments we replace the convolutional encoder and decoder with 2-layer

Original: MLP networks with 256 units each and ELU activation.
Revised: MLP networks with 256 units each and ELU activation.

Original: Additionally, the dense decoder network outputs both the mean and log standard deviation of the Gaussian distribution over observations.
Revised: Additionally, the dense decoder network outputs both the mean and log standard deviation of the Gaussian distribution over observations.

Original: The standard deviation is softly clamped between 0.1 and 1.5 as described in (Chua et al., 2018).
Revised: The standard deviation is softly clamped between 0.1 and 1.5 as described in (Chua et al., 2018).

Original: For LSP , skill vectors are 3-dimensional and are held for K = 10 steps before being updated.
Revised: For LSP , skill vectors are 3-dimensional and are held for K = 10 steps before being updated.

Original: The
Revised: The

Original: CEM method has a planning horizon of H = 10 , goes through MaxCEMiter = 4 iterations, proposes G = 16 skills and uses the top M = 4 proposals to recompute statistics in each iteration.
Revised: CEM method has a planning horizon of H = 10 , goes through MaxCEMiter = 4 iterations, proposes G = 16 skills and uses the top M = 4 proposals to recompute statistics in each iteration.

Original: The additional noise (cid:15) added to the CEM optimized distribution is Normal(0, 0.1).
Revised: The additional noise (cid:15) added to the CEM optimized distribution is Normal(0, 0.1).

Original: The backwards skill predictor shares the same architecture and settings as the feature-based decoder module described above.
Revised: The backwards skill predictor shares the same architecture and settings as the feature-based decoder module described above.

Original: It is trained with Adam with learning rate of 8 e − 5 .
Revised: It is trained with Adam with learning rate of 8 e − 5 .

Original: A. C ONTROL AS I NFERENCE
Revised: A. C ONTROL AS I NFERENCE

Original: The control as inference framework (Todorov, 2008; Levine, 2018) provides a heuristic to encourage exploration.
Revised: The control as inference framework (Todorov, 2008; Levine, 2018) provides a heuristic to encourage exploration.

Original: It performs inference in a surrogate probabilistic graphical model, where the likelihood of a trajectory being optimal (an indicator random variable O ) is a hand-designed (monotonically increasing) function of the reward, log p ( O| τ ) =
Revised: It performs inference in a surrogate probabilistic graphical model, where the likelihood of a trajectory being optimal (an indicator random variable O ) is a hand-designed (monotonically increasing) function of the reward, log p ( O| τ ) =

Original: f ( r ) .
Revised: f ( r ) .

Original: The induced surrogate posterior p ( τ |O ) places higher probability on higher reward trajectories.
Revised: The induced surrogate posterior p ( τ |O ) places higher probability on higher reward trajectories.

Original: Sampling from p ( τ |O ) is in general intractable and so a common solution is to employ variational inference for obtaining an approximation q θ ( τ ) .
Revised: Sampling from p ( τ |O ) is in general intractable and so a common solution is to employ variational inference for obtaining an approximation q θ ( τ ) .

Original: The objective is to minimize the KL -divergence with respect to the true posterior:
Revised: The objective is to minimize the KL -divergence with respect to the true posterior:

Original: [Equation]
Revised: [Equation]

Original: Note that this objective boils down to maximizing a function of the reward and an entropy term H ( · ) .
Revised: Note that this objective boils down to maximizing a function of the reward and an entropy term H ( · ) .

Original: State of the art model-free reinforcement learning approaches can be interpreted as instances of this framework (Haarnoja et al., 2018).
Revised: State of the art model-free reinforcement learning approaches can be interpreted as instances of this framework (Haarnoja et al., 2018).

Original: The connection has also been made in model-based RL (Okada & Taniguchi, 2019).
Revised: The connection has also been made in model-based RL (Okada & Taniguchi, 2019).

Original: Speciﬁcally they show how the ubiquitous MPPI/CEM planning methods can be derived from this framework when applying variational inference to simple Gaussian action distributions.
Revised: Speciﬁcally they show how the ubiquitous MPPI/CEM planning methods can be derived from this framework when applying variational inference to simple Gaussian action distributions.

Original: Though they show that more sophisticated distributions can be used in the framework as well (in principle), the only other model they investigate in the paper is a mixture Gaussians distribution.
Revised: Though they show that more sophisticated distributions can be used in the framework as well (in principle), the only other model they investigate in the paper is a mixture Gaussians distribution.

Original: A. H IERARCHICAL I NFERENCE
Revised: A. H IERARCHICAL I NFERENCE

Original: We consider hierarchical action distributions that combine amortizing low level behaviours and online planning by temporally composing these low level behaviours.
Revised: We consider hierarchical action distributions that combine amortizing low level behaviours and online planning by temporally composing these low level behaviours.

Original: Speciﬁcally we use the following variational distribution:
Revised: Speciﬁcally we use the following variational distribution:

Original: [Equation]
Revised: [Equation]

Original: Here z 1: K are latent variables deﬁning a high level plan that modulates the behaviour of the low level policy q φ ( a t | s t , z k ( t ) ) .
Revised: Here z 1: K are latent variables deﬁning a high level plan that modulates the behaviour of the low level policy q φ ( a t | s t , z k ( t ) ) .

Original: k ( t ) is an assignment deﬁning which of the K z ’s will act at time t .
Revised: k ( t ) is an assignment deﬁning which of the K z ’s will act at time t .

Original: For all our experiments, we chose a ﬁxed size window assignment
Revised: For all our experiments, we chose a ﬁxed size window assignment

Original: such that k ( t ) = (cid:98) t ( K/T ) (cid:99) .
Revised: such that k ( t ) = (cid:98) t ( K/T ) (cid:99) .

Original: This factorization corresponds to the PGM shown in 2c.
Revised: This factorization corresponds to the PGM shown in 2c.

Original: Here p ( T ) is the task distribution over which we wish to amortize.
Revised: Here p ( T ) is the task distribution over which we wish to amortize.

Original: For example using p ( s 0 ) would amortize over plans starting from different initial states.
Revised: For example using p ( s 0 ) would amortize over plans starting from different initial states.

Original: Note that the minimization over φ is outside the expectation which means that it is amortized, whereas the minimization over θ is inside such that it is inferred online.
Revised: Note that the minimization over φ is outside the expectation which means that it is amortized, whereas the minimization over θ is inside such that it is inferred online.

Original: [Equation]
Revised: [Equation]

Original: This formulation also directly relates to some meta-learning approaches that can be interpreted as hierarchical Bayes (Grant et al., 2018).
Revised: This formulation also directly relates to some meta-learning approaches that can be interpreted as hierarchical Bayes (Grant et al., 2018).

Original: There are different options for the inner and outer optimization in this type of meta-optimization.
Revised: There are different options for the inner and outer optimization in this type of meta-optimization.

Original: A. N OTE ON M UTUAL I NFORMATION
Revised: A. N OTE ON M UTUAL I NFORMATION

Original: In the paper we used the reverse skill predictor approach to estimate the mutual information objective.
Revised: In the paper we used the reverse skill predictor approach to estimate the mutual information objective.

Original: The alternate decomposition in terms of reduction in future state uncertainty is given by:
Revised: The alternate decomposition in terms of reduction in future state uncertainty is given by:

Original: [Equation]
Revised: [Equation]

Original: In this case tractable approximations of the the skill dynamics p ( s (cid:48) | z, s 0 ) and marginal dynamics p ( s (cid:48) | s 0 ) are required.
Revised: In this case tractable approximations of the the skill dynamics p ( s (cid:48) | z, s 0 ) and marginal dynamics p ( s (cid:48) | s 0 ) are required.

Original: In theory, these can be formed exactly as compositions of the policy, skill selection distribution and dynamics model:
Revised: In theory, these can be formed exactly as compositions of the policy, skill selection distribution and dynamics model:

Original: [Equation]
Revised: [Equation]

Original: The difﬁculty in using this formulation in our case is due to the marginal future state distribution p ( { s }| s 0 ) .
Revised: The difﬁculty in using this formulation in our case is due to the marginal future state distribution p ( { s }| s 0 ) .

Original: To form the marginal we rely on a ﬁnite number K monte carlo samples which is biased and our lower bound on mutual information will be reduced by log K .
Revised: To form the marginal we rely on a ﬁnite number K monte carlo samples which is biased and our lower bound on mutual information will be reduced by log K .

Original: This means that for every given initial state we wish to train on, we must also sample a sufﬁciently large number of skills from the skill selection distribution to form an acceptable marginal approximation, greatly increasing the computation cost compared to using the reverse skill predictor.
Revised: This means that for every given initial state we wish to train on, we must also sample a sufﬁciently large number of skills from the skill selection distribution to form an acceptable marginal approximation, greatly increasing the computation cost compared to using the reverse skill predictor.

Original: Function CEM( S , MaxCEMiter, G , H , K , ζ (0) ) :
Revised: Function CEM( S , MaxCEMiter, G , H , K , ζ (0) ) :

Original: Sample G initial skills { z ( g ) 1: (cid:100) H/K (cid:101) } Gg =1 from prior p ( z | ζ (0) ) for j in range(MaxCEMiter) do for g in range(G) do for each s t do τ = t //
Revised: Sample G initial skills { z ( g ) 1: (cid:100) H/K (cid:101) } Gg =1 from prior p ( z | ζ (0) ) for j in range(MaxCEMiter) do for g in range(G) do for each s t do τ = t //

Original: Imagine trajectories { ( s τ , a τ ) } t
Revised: Imagine trajectories { ( s τ , a τ ) } t

Original: + H τ = t from each s t .
Revised: + H τ = t from each s t .

Original: while τ < H , τ + + do if ( τ − t )
Revised: while τ < H , τ + + do if ( τ − t )

Original: % K == 

Revised: % K ==

Original: then
Revised: then

Original: Sample skill z 1: (cid:100) H/K (cid:101) ∼ p ( z 1: (cid:100) H/K (cid:101) | ζ ( j ) )
Revised: Sample skill z 1: (cid:100) H/K (cid:101) ∼ p ( z 1: (cid:100) H/K (cid:101) | ζ ( j ) )

Original: Sample action a τ ∼ q φ ( a τ | s τ , z 1 ) //
Revised: Sample action a τ ∼ q φ ( a τ | s τ , z 1 ) //

Original: Choose z 1 via MPC
Revised: Choose z 1 via MPC

Original: Observe new state s τ +1 and compute new intrinsic reward r
Revised: Observe new state s τ +1 and compute new intrinsic reward r

Original: i
Revised: i

Original: Compute value estimates V g from rewards of each imagined rollout
Revised: Compute value estimates V g from rewards of each imagined rollout

Original: Sort { V g } Gg =1 and use the top M sequences to update ζ ( j ) = ( µ ( j ) , Σ ( j ) )
Revised: Sort { V g } Gg =1 and use the top M sequences to update ζ ( j ) = ( µ ( j ) , Σ ( j ) )

Original: Sample G skills { z ( g ) } Gg =1 from p ( z | ζ ( j ) ) return ζ ( MaxCEMiter )
Revised: Sample G skills { z ( g ) } Gg =1 from p ( z | ζ ( j ) ) return ζ ( MaxCEMiter )