Original: I NFO B OT : T RANSFER AND E XPLORATION VIA THE I NFORMATION B OTTLENECK
Revised: I NFO B OT : T RANSFER AND E XPLORATION VIA THE I NFORMATION B OTTLENECK

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: A central challenge in reinforcement learning is discovering effective policies for tasks where rewards are sparsely distributed.
Revised: A central challenge in reinforcement learning is discovering effective policies for tasks where rewards are sparsely distributed.

Original: We postulate that in the absence of useful reward signals, an effective exploration strategy should seek out decision states .
Revised: We postulate that in the absence of useful reward signals, an effective exploration strategy should seek out decision states .

Original: These states lie at critical junctions in the state space from where the agent can transition to new, potentially unexplored regions.
Revised: These states lie at critical junctions in the state space from where the agent can transition to new, potentially unexplored regions.

Original: We propose to learn about decision states from prior experience.
Revised: We propose to learn about decision states from prior experience.

Original: By training a goal-conditioned model with an information bottleneck on the goal-dependent encoder used by the agent’s policy, we can identify decision states by examining where the model actually leverages the goal state through the bottleneck.
Revised: By training a goal-conditioned policy with an information bottleneck, we can identify decision states by examining where the model actually leverages the goal state.

Original: We ﬁnd that this simple mechanism effectively identiﬁes decision states, even in partially observed settings.
Revised: We ﬁnd that this simple mechanism effectively identiﬁes decision states, even in partially observed settings.

Original: In effect, the model learns the sensory cues that correlate with potential subgoals.
Revised: In effect, the model learns the sensory cues that correlate with potential subgoals.

Original: In new environments, this model can then identify novel subgoals for further exploration, guiding the agent through a sequence of potential decision states and through new regions of the state space.
Revised: In new environments, this model can then identify novel subgoals for further exploration, guiding the agent through a sequence of potential decision states and through new regions of the state space.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Deep reinforcement learning has enjoyed many recent successes in domains where large amounts of training time and a dense reward function are provided.
Revised: Deep reinforcement learning has enjoyed many recent successes in domains where large amounts of training time and a dense reward function are provided.

Original: However, learning to quickly perform well in environments with sparse rewards remains a major challenge.
Revised: However, learning to quickly perform well in environments with sparse rewards remains a major challenge.

Original: Providing agents with useful signals to pursue in lieu of environmental reward becomes crucial in these scenarios.
Revised: Providing agents with useful signals to pursue in lieu of environmental reward becomes crucial in these scenarios.

Original: Here, we propose to incentive agents to learn about and exploit multi-goal task structure in order to efﬁciently explore in new environments.
Revised: In this work, we propose to incentivize agents to learn about and exploit multi-goal task structure in order to efﬁciently explore in new environments.

Original: We do so by ﬁrst training agents to develop useful habits as well as the knowledge of when to break them, and then using that knowledge to efﬁciently probe new environments for reward.
Revised: We do so by ﬁrst training agents to develop useful habits as well as the knowledge of when to break them, and then using that knowledge to efﬁciently probe new environments for reward.

Original: We focus on multi-goal environments and goal-conditioned policies (Foster and Dayan, 2002; Schaul et al., 2015; Plappert et al., 2018).
Revised: We focus on multi-goal environments and goal-conditioned policies (Foster and Dayan, 2002; Schaul et al., 2015; Plappert et al., 2018).

Original: In these scenarios, a goal G is sampled from p ( G ) and the beginning of each episode and provided to the agent.
Revised: In these scenarios, a goal G is sampled from p ( G ) and the beginning of each episode and provided to the agent.

Original: The goal G provides the agent with information about the environment’s reward structure for that episode.
Revised: The goal G provides the agent with information about the environment’s reward structure for that episode.

Original: For example, in spatial navigation tasks, G might be the location or direction of a rewarding state.
Revised: For example, in spatial navigation tasks, G might be the location or direction of a rewarding state.

Original: We denote the agent’s policy π θ ( A | S, G ) , where S is the agent’s state,
Revised: We denote the agent’s policy π θ ( A | S, G ) , where S is the agent’s state,

Original: A the agent’s action, and θ the policy parameters.
Revised: A the agent’s action, and θ the policy parameters.

Original: We incentive agents to learn task structure by training policies that perform well under a variety of goals, while not overﬁtting to any individual goal.
Revised: We incentivize agents to learn task structure by training policies that perform well under a variety of goals, while not overﬁtting to any individual goal.

Original: We achieve this by training agents that, in addition to maximizing reward, minimize the policy dependence on the individual goal, quantiﬁed by the conditional mutual information
Revised: We achieve this by training agents that, in addition to maximizing reward, minimize the policy dependence on the individual goal, quantiﬁed by the conditional mutual information

Original: I ( A ; G | S ) .
Revised: I ( A ; G | S ) .

Original: Thisapproach is inspired by the information bottleneck approach (Tishby et al., 1999) to training deep neural networks for supervised learning (Alemi et al., 2017; Chalk et al., 2016; Achille and Soatto, 2016; Kolchinsky et al., 2017), where classiﬁers are trained to achieve high accuracy while simultaneously encoding as little information about the input as possible.
Revised: This is inspired by the information bottleneck approach (Tishby et al., 1999) of training deep neural networks for supervised learning (Alemi et al., 2017; Chalk et al., 2016; Achille and Soatto, 2016; Kolchinsky et al., 2017), where classiﬁers are trained to achieve high accuracy while simultaneously encoding as little information about the input as possible.

Original: This form of “information dropout” has been shown to promote generalization performance (Achille and Soatto, 2016;
Revised: This form of “information dropout” has been shown to promote generalization performance (Achille and Soatto, 2016;

Original: Alemi et al., 2017).
Revised: Alemi et al., 2017).

Original: Here, we show that minimizing goal information promotes generalization in an RL setting as well.
Revised: We show that minimizing goal information promotes generalization in an RL setting as well.

Original: We refer to this approach as InfoBot (from information bottleneck).
Revised: Our proposed model is referred as InfoBot (inspired from the Information Bottleneck framework).

Original: This approach to learning task structure can also be interpreted as encouraging agents to follow  default policy , following this default policy should result in default behaviour which  agent should follow in the absence of any additional task information (like the goal location or the relative distance to the goal or goal given in the form of language instruction etc).
Revised: This approach to learning task structure can also be interpreted as encouraging agents to follow a default policy : This is the default behaviour which the agent should follow in the absence of any additional task information (like the goal location, the relative distance to the goal or a language instruction etc).

Original: To see this, note that our regularizer can also be written as I ( A ; G | S ) =
Revised: To see this, note that our regularizer can also be written as I ( A ; G | S ) =

Original: E π θ [ D KL
Revised: E π θ [ D KL

Original: [ π θ ( A | S, G ) | π 0 ( A | S )]] , where π θ ( A | S, G ) is the agent’s multi-goal policy, E π θ denotes an expectation over trajectories generated by π θ , D KL is the Kuhlback-Leibler divergence, and π 0 ( A | S ) =
Revised: [ π θ ( A | S, G ) | π 0 ( A | S )]] , where π θ ( A | S, G ) is the agent’s multi-goal policy, E π θ denotes an expectation over trajectories generated by π θ , D KL is the Kuhlback-Leibler divergence, and π 0 ( A | S ) =

Original: (cid:80) g p ( g ) π θ ( A | S, g ) is a “default” policy with the goal marginalized out.
Revised: (cid:80) g p ( g ) π θ ( A | S, g ) is a “default” policy with the goal marginalized out.

Original: While the agent never actually follows the default policy π 0 directly, it can be viewed as what the agent might do in the absence of any knowledge about the goal.
Revised: While the agent never actually follows the default policy π 0 directly, it can be viewed as what the agent might do in the absence of any knowledge about the goal.

Original: Thus, our approach encourages the agent to learn useful behaviours and to follow those behaviours closely, except where diverting from doing so leads to signiﬁcantly higher reward.
Revised: Thus, our approach encourages the agent to learn useful behaviours and to follow those behaviours closely, except where diverting from doing so leads to signiﬁcantly higher reward.

Original: Humans too demonstrate an afﬁnity for relying on default behaviour when they can (Kool and Botvinick, 2018), which we take as encouraging support for this line of work (Hassabis et al., 2017).
Revised: Humans too demonstrate an afﬁnity for relying on default behaviour when they can (Kool and Botvinick, 2018), which we take as encouraging support for this line of work (Hassabis et al., 2017).

Original: We refer to states where diversions from default behaviour occur as decision states , based on the intuition that they require the agent not to rely on their default policy (which is goal agnostic) but instead to make a goal-dependent “decision.”
Revised: We refer to states where diversions from default behaviour occur as decision states , based on the intuition that they require the agent not to rely on their default policy (which is goal agnostic) but instead to make a goal-dependent “decision.”

Original: Our approach to exploration then involves encouraging the agent to seek out these decision states in new environments.
Revised: Our approach to exploration then involves encouraging the agent to seek out these decision states in new environments.

Original: The intuition behind this approach is that such decision states are natural subgoals for efﬁcient exploration because they are boundaries between achieving different goals (van Dijk and Polani, 2011).
Revised: Decision states are natural subgoals for efﬁcient exploration because they are boundaries between achieving different goals (van Dijk and Polani, 2011).

Original: By encouraging an agentto visit them, the agent is encouraged to 1) follow default trajectories that work across many goals (i.e could be executed in multiple different contexts) and 2) uniformly explore across the many ”branches” of decision-making.
Revised: By visiting decision states, an agent is encouraged to 1) follow default trajectories that work across many goals (i.e could be executed in multiple different contexts) and 2) uniformly explore across the many “branches” of decision-making.

Original: We encourage the visitation of decision states by ﬁrst training an agentto recognize them by training with the information regularizer introduced above.
Revised: We encourage the visitation of decision states by ﬁrst training an agent with an information regularizer to recognize decision states.

Original: Then, we freeze the agent’s policy, and use D KL [ π θ ( A | S, G ) | π 0 ( A | S )] as an exploration bonus for training a new policy.
Revised: We then freeze the agent’s policy, and use D KL [ π θ ( A | S, G ) | π 0 ( A | S )] as an exploration bonus for training a new policy.

Original: Crucially, this approach to exploration is tuned to the family of tasks the agent is trained on, and we show that it promotesmore efﬁcient exploration than other task-agnostic approaches to exploration (Houthooft et al., 2016; Pathak et al., 2017b).
Revised: Crucially, this approach to exploration is tuned to the family of tasks the agent is trained on, and we show that it promotes efﬁcient exploration than other task-agnostic approaches to exploration (Houthooft et al., 2016; Pathak et al., 2017b).

Original: To summarize our contributions:

Original: • We regularize RL agents in multi-goal settings with I ( A ; G | S ) , an approach inspired by the information bottleneck and the cognitive science of decision making, and show that it promotes generalization across tasks.

Original: • We use policies as trained above to then provide an exploration bonus for training new policies in the form of D KL [ π θ ( A | S, G ) | π 0 ( A | S )] , which encourages the agent to seek out decision states.

Original: We demonstrate that this approach to exploration performs more effectively than other state-of-the-art methods, including a count-based bonus, VIME (Houthooft et al., 2016), and curiosity (Pathak et al., 2017b).

Original: O UR APPROACH

Original: S
Revised: S

Original: G
Revised: G

Original: Z
Revised: Z

Original: A decoder encoder

Original: [Table]
Revised: [Table]

Original: +FN8eJR0ahsJoSenpqkseeD7hoxDPkDXvavePGgiFfv3vw3dmIW3MV9UPB4r4qqekGmpEHGXpyp6U8zs3Pzn0tfFhaXlssrq6cmzbWAlkhVqs8DbkDJBFooUcF5poHHgYKz4GJ/5J9dgjYyTU5wkEEn5r1ERlJwtFK3/CPr+ghXWIQghv53X0GE63t+LEN6vPnL17LXx41uucLcHVb3ajXK3Gq1XquPyFajyhpN6rlsjAqZ4LBbfvbDVOQxJCgUN6btsQw7BdcohYJhyc8NZFxc8B60LU14DKZTjL8Z0p9WCWmUalsJ0rH6fqLgsTGDOLCdMce++dcbiR957RyjZqeQSZYjJOJtUZQriikdRUNDqUGgGljChZb2Vir6XHOBNsCSDeHPp/T/5HTL9ZjrHW1XdpuTOObJN7JG1olHGmSXHJBD0iKCXJMbckfund/OrfPgPL61TjmTma/kLzhPryDVm40=</latexit>
Revised: kLzhPryDVm40=</latexit> <latexit sha1_base64="Iq/stS8ZR9f6WF3RymXwlPi7Qrs=">AAACDXicdVBNSyNBEO3R9WOzfkT3uJdeswsKMvQY8+FN8eJR0ahsJoSenpqkseeD7hoxDPkDXvavePGgiFfv3vw3dmIW3MV9UPB4r4qqekGmpEHGXpyp6U8zs3Pzn0tfFhaXlssrq6cmzbWAlkhVqs8DbkDJBFooUcF5poHHgYKz4GJ/5J9dgjYyTU5wkEEn5r1ERlJwtFK3/CPr+ghXWIQghv53X0GE63t+LEN6vPnL17LXx41uucLcHVb3ajXK3Gq1XquPyFajyhpN6rlsjAqZ4LBbfvbDVOQxJCgUN6btsQw7BdcohYJhyc8NZFxc8B60LU14DKZTjL8Z0p9WCWmUalsJ0rH6fqLgsTGDOLCdMce++dcbiR957RyjZqeQSZYjJOJtUZQriikdRUNDqUGgGljChZb2Vir6XHOBNsCSDeHPp/T/5HTL9ZjrHW1XdpuTOObJN7JG1olHGmSXHJBD0iKCXJMbckfund/OrfPgPL61TjmTma/kLzhPryDVm40=</latexit> Our contributions can be summarized as follows :

Original: <latexit sha1_base64="Iq/stS8ZR9f6WF3RymXwlPi7Qrs=">AAACDXicdVBNSyNBEO3R9WOzfkT3uJdeswsKMvQY8+FN8eJR0ahsJoSenpqkseeD7hoxDPkDXvavePGgiFfv3vw3dmIW3MV9UPB4r4qqekGmpEHGXpyp6U8zs3Pzn0tfFhaXlssrq6cmzbWAlkhVqs8DbkDJBFooUcF5poHHgYKz4GJ/5J9dgjYyTU5wkEEn5r1ERlJwtFK3/CPr+ghXWIQghv53X0GE63t+LEN6vPnL17LXx41uucLcHVb3ajXK3Gq1XquPyFajyhpN6rlsjAqZ4LBbfvbDVOQxJCgUN6btsQw7BdcohYJhyc8NZFxc8B60LU14DKZTjL8Z0p9WCWmUalsJ0rH6fqLgsTGDOLCdMce++dcbiR957RyjZqeQSZYjJOJtUZQriikdRUNDqUGgGljChZb2Vir6XHOBNsCSDeHPp/T/5HTL9ZjrHW1XdpuTOObJN7JG1olHGmSXHJBD0iKCXJMbckfund/OrfPgPL61TjmTma/kLzhPryDVm40=</latexit>
Revised: • We regularize RL agents in multi-goal settings with I ( A ; G | S ) , an approach inspired by the information bottleneck and the cognitive science of decision making, and show that it promotes generalization across tasks. • We use policies as trained above to then provide an exploration bonus for training new policies in the form of D KL [ π θ ( A | S, G ) | π 0 ( A | S )] , which encourages the agent to seek out decision states. We demonstrate that this approach to exploration performs more effectively than other state-of-the-art methods, including a count-based bonus, VIME (Houthooft et al., 2016), and curiosity (Pathak et al., 2017b).  O UR APPROACH

Original: Our objective is to train an agent on one set of tasks (environments) T ∼ p train ( T ) , but to have the agent perform well on another different, but related, set of tasks T ∼ p test ( T ) .
Revised: Our objective is to train an agent on one set of tasks (environments) T ∼ p train ( T ) , but to have the agent perform well on another different, but related, set of tasks T ∼ p test ( T ) .

Original: To do so, we propose to maximize the following objective in the training environments:
Revised: We propose to maximize the following objective in the training environments:

Original: [Equation]
Revised: [Equation]

Original: where E π θ denotes an expectation over trajectories generated by the agent’s policy, β > 0 is a tradeoff parameter, D KL is the Kuhlback-Leibler divergence, and π 0 ( A | S ) ≡ (cid:80)
Revised: where E π θ denotes an expectation over trajectories generated by the agent’s policy, β > 0 is a tradeoff parameter, D KL is the Kullback–Leibler divergence, and π 0 ( A | S ) ≡ (cid:80)

Original: g p ( g ) π θ ( A | S, g ) is a ”default” policy with the agent’s goal marginalized out.
Revised: g p ( g ) π θ ( A | S, g ) is a “default” policy with the agent’s goal marginalized out.

Original: 2. T RACTABLE BOUNDS ON INFORMATION
Revised: 2. T RACTABLE BOUNDS ON INFORMATION

Original: We parameterize the policy π θ ( A | S, G ) using an encoder p enc ( Z | S, G ) and a decoder p dec
Revised: We parameterize the policy π θ ( A | S, G ) using an encoder p enc ( Z | S, G ) and a decoder p dec

Original: ( A | S, Z ) such that π θ ( A | S, G ) =
Revised: ( A | S, Z ) such that π θ ( A | S, G ) =

Original: (cid:80) z p enc ( z | S, G ) p dec ( A | S, z ) (see ﬁgure 1).
Revised: (cid:80) z p enc ( z | S, G ) p dec ( A | S, z ) (see ﬁgure 1).

Original: The encoder output Z is meant to represent the information about the present goal G that the agent believes is important to access in the present state S in order to perform well.
Revised: The en- coder output Z is meant to represent the information about the present goal G that the agent believes is important to access in the present state S in order to perform well.

Original: The decoder takes this encoded goal information and the current state and produces a distribution over actions A .
Revised: The decoder takes this encoded goal information and the current state and produces a distribution over actions A .

Original: We suppress the dependence of p enc and p dec on θ , but θ in the union of their parameters.
Revised: We suppress the dependence of p enc and p dec on θ , but θ in the union of their parameters.

Original: Due to the data processing inequality (DPI) (Cover and Thomas, 2006), I ( Z ; G | S ) ≥ I ( A ; G | S ) , so minimizing the goal information encoded by p enc also minimizes I ( A ; G | S ) .
Revised: Due to the data processing inequality (DPI) (Cover and Thomas, 2006), I ( Z ; G | S ) ≥ I ( A ; G | S ) . Therefore, minimizing the goal information encoded by p enc also minimizes I ( A ; G | S ) .

Original: Thus, we instead maximize this lower bound on J ( θ ) :
Revised: Thus, we instead maximize this lower bound on J ( θ ) :

Original: [Equation]
Revised: [Equation]

Original: where p ( Z | S ) = (cid:80) g p ( g ) p enc ( Z | S, g ) is the marginalized encoding.
Revised: where p ( Z | S ) = (cid:80) g p ( g ) p enc ( Z | S, g ) is the marginalized encoding.

Original: In practice, performing this marginalization over the goal may often be prohibitive, since the agent might not have access to the goal distribution p ( G ) , or even if the agent does, there might be many or a continuous distribution of goals that makes the marginalization intractable.
Revised: In practice, performing this marginalization over the goal may often be prohibitive, since the agent might not have access to the goal distribution p ( G ) , or even if the agent does, there might be many or a continuous distribution of goals that makes the marginalization intractable.

Original: To avoid this marginalization, we replace p ( Z | S ) with a variational approximation q ( Z | S ) (Kingma and
Revised: To avoid this marginalization, we replace p ( Z | S ) with a variational approximation q ( Z | S ) (Kingma and

Original: Welling, 2014; Alemi et al., 2017; Houthooft et al., 2016; Strouse et al., 2018).
Revised: Welling, 2014; Alemi et al., 2017; Houthooft et al., 2016; Strouse et al., 2018).

Original: This again provides a lower bound on J ( θ ) since:
Revised: This again provides a lower bound on J ( θ ) since:

Original: [Equation]
Revised: [Equation]

Original: where the inequality in the last line, in which we replace p ( z | s ) with q ( z | s ) , follows from that
Revised: where the inequality in the last line, in which we replace p ( z | s ) with q ( z | s ) , follows from that

Original: [Equation]
Revised: [Equation]

Original: Thus, we arrive at the lower bound ˜ J ( θ ) that we maximize in practice:
Revised: Thus [Equation] we arrive at the lower bound ˜ J ( θ ) that we maximize in practice:

Original: [Equation]
Revised: [Equation]

Original: In the experiments below, we ﬁx q ( Z | S ) to be unit Gaussian, however it could also be learned, in which case its parameters should be included in θ .
Revised: In the experiments below [Equation] we ﬁx q ( Z | S ) to be unit Gaussian [Equation]

Revised: however it could also be learned

Revised: [Equation]

Revised: in which case its parameters should be included in θ .

Original: Although our approach is compatible with any
Revised: Although our approach is compatible with any

Original: RL method, we maximize ˜ J ( θ ) on-policy from sampled trajectories using a score function estimator (Williams, 1992; Sutton et al., 1999a).
Revised: RL method, we maximize ˜ J ( θ ) on-policy from sampled trajectories using a score function estimator (Williams, 1992; Sutton et al., 1999a).

Original: As derived by Strouse et al.
Revised: As derived by Strouse et al.

Original: (2018), the resulting update at time step t , which we denote ∇ θ ˜ J ( t ) , is:
Revised: [Equation] the resulting update at time step t , which we denote ∇ θ ˜ J ( t ) , is:

Original: 2. O PTIMIZATION FROM SAMPLED TRAJECTORIES

Original: [Equation]
Revised: [Equation]

Original: where ˜ R t ≡ (cid:80)
Revised: where ˜ R t ≡ (cid:80)

Original: Tu = t γ u − t ˜ r u is a modiﬁed return, ˜ r t ≡ r t + βD KL [ p enc ( Z | s t , g )
Revised: Tu = t γ u − t ˜ r u is a modiﬁed return, ˜ r t ≡ r t + βD KL [ p enc ( Z | s t , g )

Original: |
Revised: |

Original: q ( Z | s t )] is a modiﬁed reward, T is the length of the current episode, and a t , s t , and g t are the action, state, and goal at time t , respectively.
Revised: q ( Z | s t )] is a modiﬁed reward, T is the length of the current episode, and a t , s t , and g t are the action, state, and goal at time t , respectively.

Original: The ﬁrst term in the gradient comes from applying the REINFORCE update to the modiﬁed reward, and can be thought of as encouraging the agent to change the policy in the present state to revisit future states to the extent that they provide high external reward as well as low need for encoding goal information.
Revised: The ﬁrst term in the gradient comes from applying the R EINFORCE update to the modiﬁed reward, and can be thought of as encouraging the agent to change the policy in the present state to revisit future states to the extent that they provide high external reward as well as low need for encoding goal information.

Original: The second term comes from directly optimizing the policy to not rely on goal information, and can be thought of as encouraging the agent to directly alter the policy to avoid encoding goal information in the present state.
Revised: The second term comes from directly optimizing the policy to not rely on goal information, and can be thought of as encouraging the agent to directly alter the policy to avoid encoding goal information in the present state.

Original: We stress that while we take a Monte Carlo policy gradient, or REINFORCE, approach here, our regularizer and thus our method are compatible with any RL algorithm.
Revised: Note that while we take a Monte Carlo policy gradient, or R EINFORCE , approach here, our regularizer is compatible with any RL algorithm.

Revised: In practice, we estimate the marginalization over Z using a single sample throughout our experiments.

Original: 2. P OLICY AND EXPLORATION TRANSFER
Revised: 2. P OLICY AND EXPLORATION TRANSFER

Original: The end result of training with Eqn 5is that the agent learns to rely on its (goal-independent) habits as much as possible, deviating only in decision states (as introduced in Section 1) where it makes goal-dependent modiﬁcations.
Revised: By training the policy as in equation 5 the agent learns to rely on its (goal-independent) habits as much as possible, deviating only in decision states (as introduced in Section 1) where it makes goaldependent modiﬁcations.

Original: We demonstrate in Section 4 that this regularization alone already leads to generalization beneﬁts (that is, increased performance on T ∼ p test ( T ) after training on T ∼ p train ( T ) ).
Revised: We demonstrate in Section 4 that this regularization alone already leads to generalization beneﬁts (that is, increased performance on T ∼ p test ( T ) after training on T ∼ p train ( T ) ).

Original: However, our main purpose in training with Eqn 5 is to train the agent to identify decision states for the purpose of providing an exploration bonus in  new environments.
Revised: However, we train the agent to identify decision states as in equation 5, such that the learnt goaldependent policy can provide an exploration bonus in the new environments.

Original: That is, after training on
Revised: That is, after training on

Original: T ∼ p train ( T ) , we freeze the agent’s encoder p enc ( Z | S, G ) and marginal encoding q ( Z | S ) , discard the decoder p dec ( A | S, Z ) , and use the encoders to provide D KL [ p enc ( Z | S, G )
Revised: T ∼ p train ( T ) , we freeze the agent’s encoder p enc ( Z | S, G ) and marginal encoding q ( Z | S ) , discard the decoder p dec ( A | S, Z ) , and use the encoders to provide D KL [ p enc ( Z | S, G )

Original: | q ( Z | S )] as a state and goal dependent exploration bonus for training a new policy π φ ( A | S, G ) on T ∼ p test ( T ) .
Revised: | q ( Z | S )] as a state and goal dependent exploration bonus for training a new policy π φ ( A | S, G ) on T ∼ p test ( T ) .

Original: To ensure that the new agent does not pursue the exploration bonus solely (in lieu of reward), we also decay the bonus with continued visits by weighting with a count-based exploration bonus as well.
Revised: To ensure that the new agent does not pursue the exploration bonus solely (in lieu of reward), we also decay the bonus with continued visits by weighting with a count-based exploration bonus as well.

Original: That is, we divide the KL divergence by (cid:112) c ( S ) , where c ( S ) is the number of times that state has been visited during training, which is initialized to 1.
Revised: That is, we divide the KL divergence by (cid:112) c ( S ) , where c ( S ) is the number of times that state has been visited during training, which is initialized to 1.

Original: Letting r e ( t ) be the environmental reward at time t , we thus train the agent to maximize the combined reward r t :
Revised: Letting r e ( t ) be the environmental reward at time t , we thus train the agent to maximize the combined reward r t :

Original: [Equation]
Revised: [Equation]

Original: Our approach is summarized in algorithm 1.
Revised: Our approach is summarized in algorithm 1.

Original: Algorithm 1 Transfer and Exploration via the Information Bottleneck
Revised: Algorithm 1 Transfer and Exploration via the Information Bottleneck

Original: Require:
Revised: Require:

Original: [Equation]
Revised: [Equation]

Original: | S, z ) , parameterized by θ Require:
Revised: parameterized by θ Require:

Original: A variational approximation q ( Z | S ) to the goal-marginalized encoder
Revised: A variational approximation q ( Z | S ) to the goal-marginalized encoder

Original: Require:
Revised: Require:

Original: A regularization weight β
Revised: A regularization weight β

Original: Require:
Revised: Require:

Original: Another policy π φ ( A | S, G ) , along with a RL algorithm A to train it Require:
Revised: Another policy π φ ( A | S, G ) , along with a RL algorithm A to train it Require:

Original: A set of training tasks (environments) p train ( T ) and test tasks p test ( T )
Revised: A set of training tasks (environments) p train ( T ) and test tasks p test ( T )

Original: Require:
Revised: Require:

Original: A goal sampling strategy p ( G | T ) given a task T for episodes = 1 to N train do
Revised: A goal sampling strategy p ( G | T ) given a task T for episodes = 1 to N train do

Original: Sample a task T ∼ p train ( T ) and goal G ∼ p ( G | T ) Produce trajectory τ on task T with goal G using policy π θ ( A | S, G ) Update policy parameters θ over τ using Eqnend for
Revised: Sample a task T ∼ p train ( T ) and goal G ∼ p ( G | T ) Produce trajectory τ on task T with goal G using policy π θ ( A | S, G ) Update policy parameters θ over τ using Eqnend for

Original: Optional : use π θ directly on tasks sampled from p test ( T ) for episodes = 1 to N test do
Revised: Optional : use π θ directly on tasks sampled from p test ( T ) for episodes = 1 to N test do

Original: Sample a task T ∼ p test ( T ) and goal G ∼ p ( G | T ) Produce trajectory τ on task T with goal G using policy π φ ( A | S, G ) Update policy parameters φ using algorithm
Revised: Sample a task T ∼ p test ( T ) and goal G ∼ p ( G | T ) Produce trajectory τ on task T with goal G using policy π φ ( A | S, G ) Update policy parameters φ using algorithm

Original: A to maximize the reward given by Eqnend for
Revised: A to maximize the reward given by Eqnend for

Original:  R ELATED WORK
Revised:  R ELATED WORK

Original: van Dijk and Polani (2011) were the ﬁrst to point out the connection between action-goal information and the structure of decision-making.
Revised: van Dijk and Polani (2011) were the ﬁrst to point out the connection between action-goal information and the structure of decision-making.

Original: They used information to identify decision states and use them as subgoals in an options framework (Sutton et al., 1999b).
Revised: They used information to identify decision states and use them as subgoals in an options framework (Sutton et al., 1999b).

Original: We build upon their approach by combining it with deep reinforcement learning to make it more scaleable, and also modify it by using it to provide an agent with an exploration bonus, rather than subgoals for options.
Revised: We build upon their approach by combining it with deep reinforcement learning to make it more scaleable, and also modify it by using it to provide an agent with an exploration bonus, rather than subgoals for options.

Original: Our decision states are similar in spirit to the notion of ”bottleneck states” used to deﬁne subgoals in hierarchical reinforcement learning.
Revised: Our decision states are similar in spirit to the notion of ”bottleneck states” used to deﬁne subgoals in hierarchical reinforcement learning.

Original: A bottleneck state is deﬁned as one that lies on a wide variety of rewarding trajectories (McGovern and Barto, 2001; Stolle and Precup, 2002) or one that otherwise serves to separate regions of state space in a graph-theoretic sense (Menache et al., 2002; ¸Sim¸sek et al., 2005; Kazemitabar and Beigy, 2009; Machado et al., 2017).
Revised: A bottleneck state is deﬁned as one that lies on a wide variety of rewarding trajectories (McGovern and Barto, 2001; Stolle and Precup, 2002) or one that otherwise serves to separate regions of state space in a graph-theoretic sense (Menache et al., 2002; ¸Sim¸sek et al., 2005; Kazemitabar and Beigy, 2009; Machado et al., 2017).

Original: The latter deﬁnition is purely based on environmental dynamics and does not incorporate reward structure, while both deﬁnitions can lead to an unnecessary proﬁleration of subgoals.
Revised: The latter deﬁnition is purely based on environmental dynamics and does not incorporate reward structure, while both deﬁnitions can lead to an unnecessary proliferation of subgoals.

Original: To see this, consider a T-maze in which the agent starts at the bottom and two possible goals exist at either end of the top of the T. All states in this setup are bottleneck states, and hence the notion is trivial.
Revised: To see this, consider a T-maze in which the agent starts at the bottom and two possible goals exist at either end of the top of the T. All states in this setup are bottleneck states, and hence the notion is trivial.

Original: However, only the junction where the lower and upper line segments of the T meet are a decision state.
Revised: However, only the junction where the lower and upper line segments of the T meet are a decision state.

Original: Thus, we believe the notion of a decision state is a more parsimonious and accurate indication of good subgoals than is the above notions of a bottleneck state.
Revised: Thus, we believe the notion of a decision state is a more parsimonious and accurate indication of good subgoals than is the above notions of a bottleneck state.

Original: The success of our approach against state-of-the-art exploration methods (Section 4) supports this claim.
Revised: The success of our approach against state-of-the-art exploration methods (Section 4) supports this claim.

Original: We use the terminology of information bottleneck (IB) in this paper because we limit (or bottleneck) the amount of goal information used by our agent’s policy during training.
Revised: We use the terminology of information bottleneck (IB) in this paper because we limit (or bottleneck) the amount of goal information used by our agent’s policy during training.

Original: However, the correspondence is not exact: while both our method and IB limit information into the model, we maximize rewards while IB maximizes information about a target to be predicted.
Revised: However, the correspondence is not exact: while both our method and IB limit information into the model, we maximize rewards while IB maximizes information about a target to be predicted.

Original: The latter is thus a supervised learning algorithm.
Revised: The latter is thus a supervised learning algorithm.

Original: If we instead focused on imitation learning and replaced E [ r ] with I ( A ∗ ; A | S ) in Eqn 1, then our problem would correspond exactly to a variational information bottleneck (Alemi et al., 2017) between the goal G and correct action choice A ∗ (conditioned on S ).
Revised: If we instead focused on imitation learning and replaced E [ r ] with I ( A ∗ ; A | S ) in Eqn 1, then our problem would correspond exactly to a variational information bottleneck (Alemi et al., 2017) between the goal G and correct action choice A ∗ (conditioned on S ).

Original: Whye Teh et al.
Revised: Whye Teh et al.

Original: (2017) trained a policy with the same KL divergence term as in Eqn 1 for the purposes of encouraging transfer across tasks.
Revised: (2017) trained a policy with the same KL divergence term as in Eqn 1 for the purposes of encouraging transfer across tasks.

Original: They did not, however, note the connection to variational information minimization and  the  information bottleneck, nor did they leverage the learned task structure for exploration.
Revised: But this term is used completely differently context: Whye Teh et al. (2017) use a regularizer on the KL-divergence between action distributions of different policies to improve distillation, does not have any notion of goals, and is not concerned with exploration or with learning exploration strategies and transferring them to new domain. We use the variational information bottleneck, which has a KL divergence penalty on the difference between the posterior latent variable distribution and the prior.

Revised: We are not distilling multiple policies.

Original: Parallel to our work, Strouse et al.
Revised: Parallel to our work, Strouse et al.

Original: (2018) also used Eqn 1 as a training objective, however their purpose was not to show better generalization and transfer, but instead to promote the sharing and hiding of information in a multi-agent setting.
Revised: (2018) also used Eqn 1 as a training objective, however their purpose was not to show better generalization and transfer, but instead to promote the sharing and hiding of information in a multi-agent setting.

Original: Popular approaches to exploration in RL are typically based on: 1) injecting noise into action selection (e.g. epsilon-greedy, (Osband et al., 2016)), 2) encouraging “curiosity” by encouraging prediction errors of or decreased uncertainty about environmental dynamics (Schmidhuber, 1991; Houthooft et al., 2016; Pathak et al., 2017b), or 3) count-based methods which incentivize seeking out rarely visited states (Strehl and Littman, 2008; Bellemare et al., 2016; Tang et al., 2016; Ostrovski et al., 2017).
Revised: Popular approaches to exploration in RL are typically based on: 1) injecting noise into action selection (e.g. epsilon-greedy, (Osband et al., 2016)), 2) encouraging “curiosity” by encouraging prediction errors of or decreased uncertainty about environmental dynamics (Schmidhuber, 1991; Houthooft et al., 2016; Pathak et al., 2017b), or 3) count-based methods which incentivize seeking out rarely visited states (Strehl and Littman, 2008; Bellemare et al., 2016; Tang et al., 2016; Ostrovski et al., 2017).

Original: One limitation shared by all of these methods is that they have no way to leverage experience on previous tasks to improve exploration on new ones; that is, their methods of exploration are not tuned to the family of tasks the agent faces.
Revised: One limitation shared by all of these methods is that they have no way to leverage experience on previous tasks to improve exploration on new ones; that is, their methods of exploration are not tuned to the family of tasks the agent faces.

Original: Our transferrable exploration strategies approach in algorithm 1 however does exactly this.
Revised: Our transferrable exploration strategies approach in algorithm 1 however does exactly this.

Original: Another notable recent exception is Gupta et al.
Revised: Another notable recent exception is Gupta et al.

Original: (2018), which took a meta-learning approach to transferrable exploration strategies.
Revised: (2018), which took a meta-learning approach to transferable exploration strategies.

Original:  E XPERIMENTAL R ESULTS
Revised:  E XPERIMENTAL R ESULTS

Original: In this section, we demonstrate the following experimentally:
Revised: In this section, we demonstrate the following experimentally:

Original: • The policy-goal information bottleneck leads to much better policy transfer than standard RL training procedures (direct policy transfer).
Revised: • The goal-conditioned policy with information bottleneck leads to much better policy transfer than standard RL training procedures (direct policy transfer).

Original: • Using decision states as an exploration bonus leads to better performance than a variety of standard task-agnostic exploration methods (transferable exploration strategies).
Revised: • Using decision states as an exploration bonus leads to better performance than a variety of standard task-agnostic exploration methods (transferable exploration strategies).

Original: [Figure]
Revised: [Figure]

Original: 4. M INI G RID E NVIRONMENTS
Revised: 4. M INI G RID E NVIRONMENTS

Original: The ﬁrst set of environments we consider are partially observable grid worlds generated with MiniGrid (Chevalier-Boisvert and Willems, 2018), an OpenAI Gym package (Brockman et al., 2016).
Revised: The ﬁrst set of environments we consider are partially observable grid worlds generated with MiniGrid (Chevalier-Boisvert and Willems, 2018), an OpenAI Gym package (Brockman et al., 2016).

Original: We consider the MultiRoomN X S Y and a FindObjS Y task domains, as depicted in Figure 10.
Revised: We consider the MultiRoomN X S Y and a FindObjS Y task domains, as depicted in Figure 2.

Original: Both environments consist of a series of connected rooms, sometimes separated by doors that need opened.
Revised: Both environments consist of a series of connected rooms, sometimes separated by doors that need opened.

Original: In both tasks, black squares are traversible, grey squares are walls, black squares with colored borders are doors, the red triangle is the agent, and the shaded area is the agent’s visible region.
Revised: In both tasks, black squares are traversable, grey squares are walls, black squares with colored borders are doors, the red triangle is the agent, and the shaded area is the agent’s visible region.

Original: The
Revised: The

Original: MultiRoomN X S Y the environment consists of X rooms, with size at most Y , connected in random orientations.
Revised: MultiRoomN X S Y the environment consists of X rooms, with size at most Y , connected in random orientations.

Original: The agent is placed in a distal room and must navigate to a green goal square in the most distant room from the agent.
Revised: The agent is placed in a distal room and must navigate to a green goal square in the most distant room from the agent.

Original: The agent receives an egocentric view of its surrounding, consisting of 3 × 3 pixels.
Revised: The agent receives an egocentric view of its surrounding, consisting of 3 × 3 pixels.

Original: The task increases in difﬁculty with X and Y .
Revised: The task increases in difﬁculty with X and Y .

Original: The FindObjS Y environment consists of 9 connected rooms of size Y − 2 ×
Revised: The FindObjS Y environment consists of 9 connected rooms of size Y − 2 ×

Original: Y − 2 arranged in a grid.
Revised: Y − 2 arranged in a grid.

Original: The agent is placed in the center room and must navigate to an object in a randomly chosen outer room (e.g. yellow circle in bottom room in Figure 10c and blue square in top left room in Figure 10d).
Revised: The agent is placed in the center room and must navigate to an object in a randomly chosen outer room (e.g. yellow circle in bottom room in Figure 2c and blue square in top left room in Figure 2d).

Original: The agent again receives an egocentric observation, this time consisting of 7 × 7 pixels, and again the difﬁculty of the task increases with Y .
Revised: The agent again receives an egocentric observation, this time consisting of 7 × 7 pixels, and again the difﬁculty of the task increases with Y .

Original: For more details of the environment, see Appendix H.
Revised: For more details of the environment, see Appendix H.

Original: Solving these partially observable, sparsely rewarded tasks by random exploration is difﬁcult because there is a vanishing probability of reaching the goal randomly as the environments become larger.
Revised: Solving these partially observable, sparsely rewarded tasks by random exploration is difﬁcult because there is a vanishing probability of reaching the goal randomly as the environments become larger.

Original: Transferring knowledge from simpler to more complex versions of these tasks thus becomes essential.
Revised: Transferring knowledge from simpler to more complex versions of these tasks thus becomes essential.

Original: In the next two sections, we demonstrate that our approach yields 1) policies that directly transfer well from smaller to larger environments, and 2) exploration strategies that outperform other task-agnostic exploration approaches.
Revised: In the next two sections, we demonstrate that our approach yields 1) policies that directly transfer well from smaller to larger environments, and 2) exploration strategies that outperform other task-agnostic exploration approaches.

Original: 4. D IRECT P OLICY G ENERALIZATION ON M INI G RID T ASKS
Revised: 4. D IRECT P OLICY G ENERALIZATION ON M INI G RID T ASKS

Original: We ﬁrst demonstrate that training an agent with a goal bottleneck alone already leads to more effective policy transfer.
Revised: We ﬁrst demonstrate that training an agent with a goal bottleneck alone already leads to more effective policy transfer.

Original: We train policies on smaller versions of the MiniGrid environments (MultiRoomN2S and FindObjS5 and S7), but evaluate them on larger versions (MultiRoomN10S4, N10S10, and
Revised: We train policies on smaller versions of the MiniGrid environments (MultiRoomN2S and FindObjS5 and S7), but evaluate them on larger versions (MultiRoomN10S4, N10S10, and

Original: N12S10, and FindObjS7 and S10) throughout training.
Revised: N12S10, and FindObjS7 and S10) throughout training.

Original: Figure 3 compares an agent trained with a goal bottleneck (ﬁrst half of Algorithm 1) to a vanilla goal-conditioned A2C agent (Mnih et al., 2016) on MultiRoomN X S Y generalization.
Revised: Figure 3 compares an agent trained with a goal bottleneck (ﬁrst half of Algorithm 1) to a vanilla goal-conditioned A2C agent (Mnih et al., 2016) on MultiRoomN X S Y generalization.

Original: As is clear, the goal-bottlenecked agent generalizes much better.
Revised: As is clear, the goal-bottlenecked agent generalizes much better.

Original: The success rate is the number of times the agent solves a larger task with 10-12 rooms while it is being trained on a task with only 2 rooms.
Revised: The success rate is the number of times the agent solves a larger task with 10-12 rooms while it is being trained on a task with only 2 rooms.

Original: When generalizing to 10 small rooms, the agent learns to solve the task to near perfection, whereas the goal-conditioned A2C baseline only solves <50% of mazes (Figure 3a).
Revised: When generalizing to 10 small rooms, the agent learns to solve the task to near perfection, whereas the goal-conditioned A2C baseline only solves <50% of mazes (Figure 3a).

Original: Table 1 compares the same two agents on FindObjS Y generalization.
Revised: Table 1 compares the same two agents on FindObjS Y generalization.

Original: In addition, this comparison includes an ablated version of our agent with β = 0 , that is an agent with the same architecture as in Figure 1 but with the no information term in its training objective.
Revised: In addition, this comparison includes an ablated version of our agent with β = 0 , that is an agent with the same architecture as in Figure 1 but with the no information term in its training objective.

Original: This is to ensure that our method’s success is not due to the architecture alone.
Revised: This is to ensure that our method’s success is not due to the architecture alone.

Original: As is evident, the goal-bottlenecked agent again generalizes much better.
Revised: As is evident, the goal-bottlenecked agent again generalizes much better.

Original: We analyzed the agent’s behaviour to understand the intuition of why it generalizes well.
Revised: We analyzed the agent’s behaviour to understand the intuition of why it generalizes well.

Original: In the
Revised: In the

Original: MultiRoomN X S Y environments, we ﬁnd that the agent quickly discovers a wall following strategy.
Revised: MultiRoomN X S Y environments, we ﬁnd that the agent quickly discovers a wall following strategy.

Original: Since these environments are partially observable, this is indeed a good strategy that also generalizes well to larger mazes.
Revised: Since these environments are partially observable, this is indeed a good strategy that also generalizes well to larger mazes.

Original: In the FindObjS Y environments, on the other hand, the agent sticks toward the center of the rooms, making a beeline from doorway to doorway.
Revised: In the FindObjS Y environments, on the other hand, the agent sticks toward the center of the rooms, making a beeline from doorway to doorway.

Original: This is again a good strategy, because the agent’s ﬁeld of view in these experiments is large enough to see the entire room in which its in to determine if the goal object is present or not.
Revised: This is again a good strategy, because the agent’s ﬁeld of view in these experiments is large enough to see the entire room in which its in to determine if the goal object is present or not.

Original: [Table]
Revised: [Table]

Original: [-Timesteps × 
 S u cce ss-]
Revised: [-Timesteps × 
 S u cce ss-]on FindObjS Y .

Original: R a t e ( % ) Trained on N2S6, Evaluation on[-N10S InfoBotBaseline (a) 0 .0 .0 .0 .1 . 
 Timesteps × 
-]
Revised: Agents trained on FindObjS5, and evaluated on[-N10S InfoBotBaseline (a) 0 .0 .0 .0 .1 . 
 Timesteps × 
-]

Original: S u cce ss
Revised: FindObjS7 and S10.

Original: R a t e ( % )
Revised: [Figure]

Original: [-Trained on N2S6, Evaluation on N10S InfoBotBaseline (b) 0 .0 .0 .0 .1 . 
 Timesteps × 10 6 − S u cce ss-]
Revised: [-Trained on N2S6, Evaluation on N10S InfoBotBaseline (b) 0 .0 .0 .0 .1 . 
 Timesteps × 10 6 − S u cce ss-][Figure]

Original: R a t e ( % )
Revised: t e ( [Figure]

Original: Trained

Original: [Figure]
Revised: [Figure]

Original: N2S6, Evaluation on N12S
Revised: t

Original: InfoBotBaseline (
Revised: [Figure] [Figure] [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: 4. T RANSFERABLE E XPLORATION S TRATEGIES ON M INI G RID T ASKS
Revised: 4. T RANSFERABLE E XPLORATION S TRATEGIES ON M INI G RID T ASKS

Original: We now evaluate our approach to exploration (the second half of Algorithm 1).
Revised: We now evaluate our approach to exploration (the second half of Algorithm 1).

Original: We train agents with a goal bottleneck on one set of environments (MultiRoomN2S6) where they learn the sensory cues that correspond to decision states.
Revised: We train agents with a goal bottleneck on one set of environments (MultiRoomN2S6) where they learn the sensory cues that correspond to decision states.

Original: Then, we that knowledge to guide exploration one another set of environments (MultiRoomN3S4, MultiRoomN4S4, and MultiRoomN5S4).We compare to several standard task-agnostic exploration methods, including count-based exploration (Eqn 6 without the D KL , that a bonus of β/ (cid:112) c ( s ) ), VIME (Houthooft et al., 2016), and curiosity-driven exploration (Pathak et al., 2017b), as well as a goal-conditioned A2C baseline with no exploration bonuses.
Revised: We then use the identiﬁed decision states to guide exploration on another set of environments (MultiRoomN3S4, MultiRoomN4S4, and MultiRoomN5S4). We compare to several standard task-agnostic exploration methods, including count-based exploration (Eqn 6 without the D KL , that a bonus of β/ (cid:112) c ( s ) ), VIME (Houthooft et al., 2016), and curiosity-driven exploration (Pathak et al., 2017b), as well as a goal-conditioned A2C baseline with no exploration bonuses.

Original: Results are shown in Table 5 and Figure 4.
Revised: Results are shown in Table 2 and Figure 4.

Original: On a maze with three rooms, the count-based method and curiosity-driven exploration slightly outperform the proposed learned exploration strategy.
Revised: On a maze with three rooms, the count-based method and curiosity-driven exploration slightly outperform the proposed learned exploration strategy.

Original: However, as the number of rooms increases, the count-based method and VIME fail completely and the curiosity-based method degrades to only 54% success rate.
Revised: However, as the number of rooms increases, the count-based method and VIME fail completely and the curiosity-based method degrades to only 54% success rate.

Original: This is in contrast to the proposed exploration strategy, which by learning the structure of the task, maintains a high success rate of 85%.[Figure]
Revised: This is in contrast to the proposed exploration strategy, which by learning the structure of the task, maintains a high success rate of 85%.

Original: A v e r a g e T a s k R e t u r n
Revised: 4. G OAL - BASED N AVIGATION T ASKS

Original: [Figure] [Figure]
Revised: In this task, we use a partially observed goal based MiniPacMan environment as shown in Figure 5.

Original: A
Revised: The agent navigates a maze, and tries to reach the goal. The agent sees only a partial window around itself.

Original: [Figure] e [Figure] g e
Revised: The agent only gets a reward of “1” when it reaches the goal.

Original: T a 
Revised: For standard RL algorithms, these tasks are difﬁcult to solve due to the partial observability of the environment, sparse reward (as the agent receives a reward only after reaching the goal), and low probability of reaching the goal via random walks (precisely because these junction states are crucial states where the right action must be taken and several junctions need to be crossed).

Original: [Figure] u [Figure]
Revised: This environment is more challenging as compared to the Minigrid environment, as this environment also has dead ends as well as more complex branching.

Original: n
Revised: We ﬁrst demonstrate that training an agent with a goal bottleneck alone leads to more effective policy transfer.

Original: [Figure] [Figure]
Revised: We train policies on smaller versions of this goal based MiniPacMan environment environments (6 x 6 maze), but evaluate them on larger versions (11 X 11) throughout training.

Original: Timesteps × 100 .0 .0 .0 .0 .1 .
Revised: Timesteps × 100 .0 .0 .0 .0 .1 .

Original: A v e r a g e
Revised: A v e r a g e

Original: T a s k R e t u r n
Revised: T a s k R e t u r n

Original: Exploration 
Revised: Exploration Bonus on N3S

Original: [Figure] (c) 4. T RANSFERABLE E XPLORATION S TRATEGIES C ONTINUOUS C ONTROL
Revised: InfoBotCount-BasedBaseline (a) 0 . 0 0 .

Original: [Figure] To show that the InfoBot architecture can also be applied to continuous control, we evaluated the performance of InfoBot on three continuous control tasks from OpenAI Gym (Brockman et al., 2016). Because InfoBot depends on the goal, in the control domains, we use high value states as an approximation to the goal state following Goyal et al. We maintain a buffer of high value states, and at each update, we sample a high value state from the buffer which acts as a proxy for the goal.
Revised: Timesteps × 10 4 − 0 .0 .0 .0 .0 .0 .1 .

Original: We compared to proximal policy optimization (PPO) (Schulman et al., 2017), as well as two ablated versions of our model: 1) instead of taking high value states, we take low value states from the buffer as proxy to the goal (”InfoBot-low-value”) and 2) we use the same InfoBot policy architecture but do not use the information regularizer (i.e. β = 0 ) (”InfoBot-zero-kl”). The results in Figure 8 show that InfoBot improves over all three alternatives. 4. T RANSFERABLE E XPLORATION S TRATEGIES FOR A TARI
Revised: A v e r a g e

Original: We further evaluate our expeirments on few Atari games (Bellemare et al., 2013) using A2C and compare it with our proposed InfoBot framework. In this experiment, our goal is to show that InfoBot can generalize to even more complex domains compared to the maze tasks above. As in the control experiments, we again use high value states as a proxy to the goal Goyal et al. and maintain a buffer of 20000 states to sample and prioritize high value states.
Revised: T a s k R e t u r n

Original: We evaluate our proposed model on 4 Atari games (Pong, QBert, Seaquest and BreakOut) as shown in ﬁgure 6. We ﬁnd that compared to a vanilla A2C agent, our InfoBot A2C agent performs signiﬁcantly better on [-these Atari tasks. Our experiments are averaged over three random seeds. 5 Timesteps × 10− 20 −
-]
Revised: Exploration Bonus on [-these Atari tasks. Our experiments are averaged over three random seeds. 5 Timesteps × 10− 20 −
-]N4S

Original: A v e r a g e T a s k R e t u r n Pong InfoBotBaseline 0 . 0 0  . {+
+}
Revised: InfoBotCount-BasedBaseline (b) 0 .0 .1 .1 .2 .2 .3 . {+
+}

Original: 5 Timesteps × [-10 
 A v e r a g e T a s k R e t u r n Qbert-]
Revised: Timesteps × [-10 
 A v e r a g e T a s k R e t u r n Qbert-]100 .0 .0 .0 .0 .1 .

Original: [-InfoBotBaseline     5 Timesteps × 
-]
Revised: [-InfoBotBaseline     5 Timesteps × 
-]

Original: A v e r a g e T a s k R e t u r n
Revised: A v e r a g e T a s k R e t u r n

Original: Seaquest InfoBotBaseline     5 Timesteps × 10 6 − 2 .0 . A v e r a g e T a s k
Revised: Exploration Bonus on N5S

Original: R e t u r n BreakOut InfoBotBaseline
Revised: InfoBotCount-BasedBaseline (c)

Original: [Table]
Revised: [Table]

Original: Here, we use a partially observed goal based MiniPacMan environment as shown in Figure 7. The agent navigates a maze, and tries to reach the goal. The agent sees only a partial window around itself.
Revised: maze

Original: The agent only gets a reward of “1” when it reaches the goal. For standard RL algorithms, these tasks are difﬁcult to solve due to the partial observability of the environment, sparse reward (as the agent receives a reward only after reaching the goal), and low probability of reaching the goal via random walks (precisely because these junction states are crucial states where the right action must be taken and several junctions need to be crossed). This environment is more challenging as compared to the
Revised: Comparison of our proposed method to regular actor-critic methods, UVFA and other hierarchical approaches.

Original: Minigrid environment, as this environment also has dead ends as well as more complex branching. We ﬁrst demonstrate that training an agentwith a goal bottleneck alone leads to more effective policy transfer. We train policies on smaller versions of this goalbased MiniPacMan environment environments ( 6 x 6 maze), but evaluate them on larger versions (11 X 11) throughout training.
Revised: Results shown times agent goal

Original: Table.
Revised: Table.

Original: 3 compares an agent trained with a goal bottleneck (ﬁrst half of Algorithm 1) to a vanilla goal-conditioned A2C agent (Mnih et al., 2016), exploration methods like count-based exploration, curiosity driven exploration (Pathak et al., 2017a) and Feudal RL agorithm (Vezhnevets et al., 2017).
Revised: 3 compares an agent trained with a goal bottleneck (ﬁrst half of Algorithm 1) to a vanilla goal-conditioned A2C agent (Mnih et al., 2016), exploration methods like count-based exploration, curiosity driven exploration (Pathak et al., 2017a) and Feudal RL algorithm (Vezhnevets et al., 2017).

Original: The goal-bottlenecked agent generalizes much better.
Revised: The goal-bottlenecked agent generalizes much better.

Original: The success rate is the number of times the agent solves a larger task while it is being trained on a task with only 2 rooms.
Revised: The success rate is the number of times the agent solves a larger task while it is being trained on a task with only 2 rooms.

Original: When generalizing to larger maze, the agent learns to solve the task 64% of the times, whereas other agents solve <50% of mazes (Table 3).
Revised: When generalizing to larger maze, the agent learns to solve the task 64% of the times, whereas other agents solve <50% of mazes (Table 3).

Original: For representing the goals, we experiment with 2 versions.
Revised: For representing the goals, we experiment with 2 versions.

Original: 1) In which we give the agent’s relative distance to the goal as the goal representation.
Revised: 1) In which we give the agent’s relative distance to the goal as the goal representation.

Original: In which we give the top down image of the goal.
Revised: In which we give the top down image of the goal.

Original: For our experiments, the baseline in which we give the relative distance to the goal worked better.
Revised: For our experiments, the baseline in which we give the relative distance to the goal worked better.

Original: 4. M IDDLE GROUND BETWEEN M ODEL B ASED RL AND M ODEL FREE RL
Revised: 4. M IDDLE GROUND BETWEEN M ODEL B ASED RL AND M ODEL FREE RL

Original: We further demonstrate the idea of decision states in a planning goal-based navigation task that uses a combination of model-based and model-free RL.
Revised: We further demonstrate the idea of decision states in a planning goal-based navigation task that uses a combination of model-based and model-free RL.

Original: Identifying useful decision states can provide a comfortable middle ground between model-free reasoning and model-based planning.
Revised: Identifying useful decision states can provide a comfortable middle ground between model-free reasoning and model-based planning.

Original: For example, imagine planning over individual decision states, while using model-free knowledge to navigate between bottlenecks: aspects of the environment that are physically complex but vary little between problem instances are handled in a model-free way (the navigation between decision points), while the particular decision points that are relevant to the task can be handled by explicitly reasoning about causality, in a more abstract representation.
Revised: For example, imagine planning over individual decision states, while using model-free knowledge to navigate between bottlenecks: aspects of the environment that are physically complex but vary little between problem instances are handled in a model-free way (the navigation between decision points), while the particular decision points that are relevant to the task can be handled by explicitly reasoning about causality, in a more abstract representation.

Original: We demonstrate this using a similar setup as in imagination augmented agents (Weber et al., 2017).
Revised: We demonstrate this using a similar setup as in imagination augmented agents (Weber et al., 2017).

Original: In imagination augmented agents, model free agents are augmented with imagination, such that the imagination model can be queried to make predictions about the future.
Revised: In imagination augmented agents, model free agents are augmented with imagination, such that the imagination model can be queried to make predictions about the future.

Original: We use the dynamics models to simulate imagined trajectories, which are then summarized by a neural network and this summary is provided as additional context to a policy network.
Revised: We use the dynamics models to simulate imagined trajectories, which are then summarized by a neural network and this summary is provided as additional context to a policy network.

Original: Here, we use the output of the imagination module as a “goal” and we want to show that only near the decision points (i e potential subgoals) the agent wants to make use of the information which is a result of running imagination module. 
Revised: We use the output of the imagination module as a “goal” and we want to show that only near the decision points (i.e potential subgoals) the agent wants to make use of the information which is a result of running imagination module. [Figure]

Original: [Figure]
Revised: [Figure]

Revised:  C ONNECTIONS TO N EUROSCIENCE AND C OGNITIVE S CIENCE

Original: [Figure]
Revised: The work we have presented bears some interesting connections to cognitive science and neuroscience. Both of these ﬁelds draw a fundamental distinction between automatic and controlled action selection (Miller and Cohen, 2001). In automatic responses, perceptual inputs directly trigger actions according to a set of habitual stimulus-response associations. In controlled behaviour, automatic responses are overridden in order to align behaviour with a more complete representation of task context, including current goals.

Original: C ONCLUSION
Revised: As an example, on the drive to work, automatic responding may trigger the appropriate turn at a familiar intersection but top-down control may be needed to override this response if the same intersection is encountered on the route to a less routine destination. As is readily evident our InfoBot architecture contains two pathways that correspond rather directly to the automatic and controlled pathways that have been posited in cognitive neuroscience models (Miller and Cohen, 2001). In the neuroscience context, representation of task context and the function of overriding automatic responses has been widely linked with the prefrontal cortex (Miller and Cohen, 2001), and it is interesting to consider the route within InfoBot from goal to action representations in this light. Notably, recent work has suggested that prefrontal control processes are associated with subjective costs; ceteris paribus, human decision-makers will opt for habitual or automatic routes to behaviour.

Revised: This of course aligns with InfoBot, and in particular with the KL term in Equation 1.

Revised: This correspondence with neuroscience provides some indirect encouragement for the approach implemented in the present work.

Revised: In turn, the InfoBot framework provides an indication for why a cost of control may exist in human cognition, namely that this encourages the emergence of useful habits, with payoffs for efﬁcient exploration and transfer.

Revised:  C ONCLUSION

Original: In this paper we proposed to train agents to develop “default behaviours” as well as the knowledge of when to break those behaviour, using an information bottleneck between the agent’s goal and policy.
Revised: In this paper, we proposed to train agents to develop “default behaviours” as well as the knowledge of when to break those behaviour, using an information bottleneck between the agent’s goal and policy.

Original: We demonstrated empirically that this training procedure leads to better direct policy transfer across tasks.
Revised: We demonstrated empirically that this training procedure leads to better direct policy transfer across tasks.

Original: We also demonstrated that the states in which the agent learns to deviate from its habits, which we call ”decision states”, can be used as the basis for a learned exploration bonus that leads to more effective training than other task-agnostic exploration methods. 
Revised: We also demonstrated that the states in which the agent learns to deviate from its habits, which we call ”decision states”, can be used as the basis for a learned exploration bonus that leads to more effective training than other task-agnostic exploration methods. A CKNOWLEDGEMENTS

Revised: The authors acknowledge the important role played by their colleagues at Mila throughout the duration of this work.

Revised: AG would like to thank Doina Precup, Maxime Chevalier-Boisvert and Alexander Neitz for very useful discussions.

Revised: AG would like to thank Konrad Kording, Rosemary Nan Ke, Jonathan

Revised: Binas, Bernhard Schoelkopf for useful discussions.

Original: A M ATHEMATICAL F RAMEWORK
Revised: The authors would also like to thank Yee Whye Teh, Peter Henderson, Emmanuel Bengio, Philip Bachman, Michael Noukhovitch for feedback on the draft. The authors are grateful to NSERC, CIFAR, Google, Samsung, Nuance, IBM, Canada Research Chairs, Canada Graduate Scholarship Program, Nvidia for funding, and Compute Canada for computing resources.

Revised: We are very grateful to Google for giving Google Cloud credits used in this project.

Revised: A M ATHEMATICAL F RAMEWORK

Original: Here, we show that the proposed approach is equivalent to regularizing agents with a variational upper bound on the mutual information between goals and actions given states.
Revised: We show that the proposed approach is equivalent to regularizing agents with a variational upper bound on the mutual information between goals and actions given states.

Original: Recall that the variational information bottleneck objective (Alemi et al., 2016; Tishby et al., 2000) is formulated as the maximization of I ( Z, Y ) − βI ( Z, X ) .
Revised: Recall that the variational information bottleneck objective (Alemi et al., 2016; Tishby et al., 2000) is formulated as the maximization of I ( Z, Y ) − βI ( Z, X ) .

Original: In our setting, the input ( X ) corresponds to the goal of the agent ( G ) and ( A ) corresponds to the target output.
Revised: In our setting, the input ( X ) corresponds to the goal of the agent ( G ) and ( A ) corresponds to the target output.

Original: We assume that the joint distribution p ( G, A, Z | S ) factorizes as follows: p ( G, A, Z | S )
Revised: We assume that the joint distribution p ( G, A, Z | S ) factorizes as follows: p ( G, A, Z | S )

Original: =
Revised: =

Original: p ( Z | G, A, S ) p
Revised: p ( Z | G, A, S ) p

Original: ( A | Z, S ) p ( G | S )
Revised: ( A | Z, S ) p ( G | S )

Original: = p ( Z | G, S ) p ( A | Z, S ) p ( G | S )
Revised: = p ( Z | G, S ) p ( A | Z, S ) p ( G | S )

Original: i.e., we assume p ( Z | G, A, S )
Revised: i.e., we assume p ( Z | G, A, S )

Original: = p ( Z | G, S ) , corresponding to the Markov chain G
Revised: = p ( Z | G, S ) , corresponding to the Markov chain G

Original: → Z → A .
Revised: → Z → A .

Original: The Data Processing Inequality (DPI) (Beaudry and Renner, 2012; Kinney and Atwal, 2014; Achille and Soatto, 2017) for a Markov chain x
Revised: The Data Processing Inequality (DPI) (Beaudry and Renner, 2012; Kinney and Atwal, 2014; Achille and Soatto, 2017) for a Markov chain x

Original: → z → y ensures that I ( x ; z ) ≥ I ( x ; y ) .
Revised: → z → y ensures that I ( x ; z ) ≥ I ( x ; y ) .

Original: Hence for Infobot, it implies,
Revised: Hence for Infobot, it implies,

Original: [Equation]
Revised: [Equation]

Original: To get an upper bound on I ( G ; Z | S ) , we ﬁrst would get an upper bound on I ( G ; Z | S = s ) , and then we average over p ( s ) to get the required upper bound.
Revised: To get an upper bound on I ( G ; Z | S ) , we ﬁrst would get an upper bound on I ( G ; Z | S = s ) , and then we average over p ( s ) to get the required upper bound.

Original: We get the following result,
Revised: We get the following result,

Original: [Equation]
Revised: [Equation]

Original: Here, we assume that marginalizing over goals to get p ( z | s ) =
Revised: Here, we assume that marginalizing over goals to get p ( z | s ) =

Original: (cid:80) g p ( g ) p ( z | s, g ) is intractable, and so we approximate it with a normal prior p prior =
Revised: (cid:80) g p ( g ) p ( z | s, g ) is intractable, and so we approximate it with a normal prior p prior =

Original: N (0 , 1) .
Revised: N (0 , 1) .

Original: Since the cross-entropy between p ( z | s ) and p prior ( z ) is larger than between p ( z | s ) and itself, we get the following upper bound:
Revised: Since the cross-entropy between p ( z | s ) and p prior ( z ) is larger than between p ( z | s ) and itself, we get the following upper bound:

Original: [Equation]
Revised: [Equation]

Original: Averaging over state probabilities gives
Revised: Averaging over state probabilities gives

Original: [Equation]
Revised: [Equation]

Original: Using Eq.
Revised: Using Eq.

Original: 7 , we can get an upper bound on the mutual information between immediate action and goal.
Revised: 7 , we can get an upper bound on the mutual information between immediate action and goal.

Original: Here r ( z ) is a ﬁxed prior.
Revised: Here r ( z ) is a ﬁxed prior.

Original: [Equation]
Revised: [Equation]

Original: Hence, minimizing the KL between p prior ( z ) and p ( z | s, g ) penalizes the use of information about the goal by the policy, so that when the policy decides to use information about the goal, it mustbe worthwhile, is supposed to a behaviour.
Revised: Hence, minimizing the KL between p prior ( z ) and p ( z | s, g ) penalizes the use of information about the goal by the policy, so that when the policy decides to use information about the goal, it must

Original: 6 Timesteps × 10−−
Revised: [Figure] [Figure]

Original: A v e r a g e T a s k R e t u r n
Revised: A v e r a g e T a s k R e t u r n

Original: Transfer to Pong
Revised: [Figure] A. T RANSFERABLE E XPLORATION S TRATEGIES FOR C ONTINUOUS C ONTROL To show that the InfoBot architecture can also be applied to continuous control, we evaluated the performance of InfoBot on three continuous control tasks from OpenAI Gym (Brockman et al., 2016).

Original: InfoBotBaseline (a) Transfer task to Pong 0 . 0 0 .
Revised: Because InfoBot depends on the goal, in the control domains, we use high value states as an approximation to the goal state following Goyal et al.

Original: [-5 Timesteps × 
-]
Revised: [-5 Timesteps × 
-]We maintain a buffer of high value states, and at each update, we sample a high value state from the buffer which acts as a proxy for the goal. We compared to proximal policy optimization (PPO) (Schulman et al., 2017), as well as two ablated versions of our model: 1) instead of taking high value states, we take low value states from the buffer as proxy to the goal (”InfoBot-low-value”) and 2) we use the same InfoBot policy architecture but do not use the information regularizer (i.e. β = 0 ) (”InfoBot-zero-kl”). The results in Figure 7 show that InfoBot improves over all three alternatives.

Original: A v e r a g e T a s k R e t u r n Transfer to Qbert
Revised: A. T RANSFERABLE E XPLORATION S TRATEGIES FOR A TARI

Original: InfoBotBaseline
Revised: We further evaluate our experiments on few Atari games (Bellemare et al., 2013) using A2C and compare it with our proposed InfoBot framework. In this experiment, our goal is to show that InfoBot can generalize to even more complex domains compared to the maze tasks above. As in the control experiments, we again use high value states as a proxy to the goal Goyal et al. and maintain a buffer of 20000 states to sample and prioritize high value states.

Original: [Figure]
Revised: We evaluate our proposed model on 4 Atari games (Pong, QBert, Seaquest and BreakOut) as shown in ﬁgure 8. We ﬁnd that compared to a vanilla A2C agent, our InfoBot A2C agent performs signiﬁcantly better on these Atari tasks. Our experiments are averaged over three random seeds.

Original: B T RANSFERRED E XPLORATION S TRATEGY IN A TARI
Revised: B T RANSFERRED E XPLORATION S TRATEGY IN A TARI

Original: We demonstrate that the encoder can be used to transfera exploration strategy across Atari games to help agents learn a new game quickly.
Revised: We demonstrate that the encoder can be used to transfera exploration strategy across Atari games to help agents learn a new game quickly.

Original: To initialize the encoder, we train an agent on the game
Revised: To initialize the encoder, we train an agent on the game

Original: Seaquest, where the agent is trained to identify the decision states.
Revised: Seaquest, where the agent is trained to identify the decision states.

Original: We then re-use this encoder on another Atari environment to provide an exploration bonus.
Revised: We then re-use this encoder on another Atari environment to provide an exploration bonus.

Original: Our experiments on the Atari games are evaluated on the games of Pong and Qbert.
Revised: Our experiments on the Atari games are evaluated on the games of Pong and Qbert.

Original: On Pong, an agent with learns to get a task return of 20 in 3M time steps.
Revised: On Pong, an agent with learns to get a task return of 20 in 3M time steps.

Original: We show that the identiﬁcation of bottleneck states can be used to transfer exploration strategies across Atari games (Fig. 9). 
Revised: We show that the identiﬁcation of bottleneck states can be used to transfer exploration strategies across Atari games (Fig. 9). 6 Timesteps × 10−− A v e r a g e T a s k R e t u r n

Revised: Transfer to Pong

Revised: InfoBotBaseline (a) Transfer task to Pong 0 . 0 0 .

Original: C{+5 Timesteps × 
+} A LGORITHM I MPLEMENTATION D ETAILS
Revised: {+5 Timesteps × 
+} A v e r a g e T a s k R e t u r n Transfer to Qbert

Original:  We evaluate the InfoBot framework using Adavantage Actor-Critic (A2C) to learn a policy π θ ( a | s, g ) conditioned on the goal.
Revised: InfoBotBaseline [Figure] C A LGORITHM I MPLEMENTATION D ETAILS We evaluate the InfoBot framework using Adavantage Actor-Critic (A2C) to learn a policy π θ ( a | s, g ) conditioned on the goal.

Original: To evaluate the performance of InfoBot, we use a range of maze multiroom tasks from the gym-minigrid framework (Chevalier-Boisvert and Willems, 2018) and the A2C implementation from (Chevalier-Boisvert and Willems, 2018).
Revised: To evaluate the performance of InfoBot, we use a range of maze multiroom tasks from the gym-minigrid framework (Chevalier-Boisvert and Willems, 2018) and the A2C implementation from (Chevalier-Boisvert and Willems, 2018).

Original: For the maze tasks, we used agent’s relative distance to the absolute goal position as "goal".
Revised: For the maze tasks, we used agent’s relative distance to the absolute goal position as "goal".

Original: For the maze environments, we use A2C with 48 parallel workers.
Revised: For the maze environments, we use A2C with 48 parallel workers.

Original: Our actor network and critic networks consist of two and three fully connected layers respectively, each of which have 128 hidden units.
Revised: Our actor network and critic networks consist of two and three fully connected layers respectively, each of which have 128 hidden units.

Original: The encoder network is also parameterized as a neural network, which consists of 1 fully connected layer.
Revised: The encoder network is also parameterized as a neural network, which consists of 1 fully connected layer.

Original: We use RMSProp with an initial learning rate of 0 . 0007 to train the models, for both InfoBot and the baseline for a fair comparison.
Revised: We use RMSProp with an initial learning rate of 0 . 0007 to train the models, for both InfoBot and the baseline for a fair comparison.

Original: Due to the partially observable nature of the environment, we further use a LSTM to encode the state and summarize the past observations.
Revised: Due to the partially observable nature of the environment, we further use a LSTM to encode the state and summarize the past observations.

Original: For the Atari experiments, we use the open-source A2C implementation from Kostrikov (2018) and condition the goal state into the policy network.
Revised: For the Atari experiments, we use the open-source A2C implementation from Kostrikov (2018) and condition the goal state into the policy network.

Original: For the actor network, we use 3 convolution layers with ReLU activations.
Revised: For the actor network, we use 3 convolution layers with ReLU activations.

Original: For training the networks, we use RMSProp with pre-deﬁned learning rates for our algorithm and the baseline.
Revised: For training the networks, we use RMSProp with pre-deﬁned learning rates for our algorithm and the baseline.

Original: The goal state in our experiments is used as the high value state following (Goyal et al., 2018).
Revised: The goal state in our experiments is used as the high value state following (Goyal et al., 2018).

Original: For the Mujoco Experiments, we use Proximal Policy Optimization (PPO)
Revised: For the Mujoco Experiments, we use Proximal Policy Optimization (PPO)

Original: (Schulman et al., 2017) with the open-source implementation available in Kostrikov (2018), using the deﬁned architectures and hyperparameters as given in the repository for both our algorithm and the baseline.
Revised: (Schulman et al., 2017) with the open-source implementation available in Kostrikov (2018), using the deﬁned architectures and hyperparameters as given in the repository for both our algorithm and the baseline.

Original: Again, for the goal-conditioned policies, the goal state for Mujoco experiments is deﬁned as the high value state following (Goyal et al., 2018).
Revised: Again, for the goal-conditioned policies, the goal state for Mujoco experiments is deﬁned as the high value state following (Goyal et al., 2018).

Original: For reproducibility purposes of our experiments, we will further release the code on github that will be available on
Revised: For reproducibility purposes of our experiments, we will further release the code on github that will be available on

Original: Hyperparameter Search for Grid World Experiments :
Revised: Hyperparameter Search for Grid World Experiments :

Original: For the proposed method, we only varied the weight of KL loss.
Revised: For the proposed method, we only varied the weight of KL loss.

Original: We tried 3 different values 0.1, 0.9, 0.01, 0.09, 0.005 for each of the env. and[Table] plotted the results which gave the most improvement.
Revised: We tried 5 different values 0.1, 0.9, 0.01, 0.09, 0.005 for each of the env. and plotted the results which gave the most improvement.

Original: We used CNN policy for the case of FindObjSY env, and MLP policy for the case of MultiRoomNXSY.
Revised: We used CNN policy for the case of FindObjSY env, and MLP policy for the case of MultiRoomNXSY.

Original: We used the agent’s relative distance to goal as a goal for our goal conditioned policy.
Revised: We used the agent’s relative distance to goal as a goal for our goal conditioned policy.

Original: Hyperparameter Search for Control experiments We have not done any hyper-parameter search for the baseline.
Revised: Hyperparameter Search for Control experiments We have not done any hyper-parameter search for the baseline.

Original: For the proposed method, we only varied the weight of KL loss.
Revised: For the proposed method, we only varied the weight of KL loss.

Original: We tried 3 different values 0.5, 0.01, 0.09 for each of the env. and plotted the results which gave the most improvement over the PPO baseline. 
Revised: We tried 3 different values 0.5, 0.01, 0.09 for each of the env. and plotted the results which gave the most improvement over the PPO baseline. https:// github com

Revised: /anonymous

Revised: [Table]

Original: D M ULTIAGENT C OMMUNICATION
Revised: D M ULTIAGENT C OMMUNICATION

Original: Here, we want to show that by training agents to develop “default behaviours” as well as the knowledge of when to break those behaviours, using an information bottleneck can also help in other scenarios like multi-agent communication.
Revised: Here, we want to show that by training agents to develop “default behaviours” as well as the knowledge of when to break those behaviours, using an information bottleneck can also help in other scenarios like multi-agent communication.

Original: Consider multiagent communication, where in order to solve a task, agents require communicating with another agents.
Revised: Consider multiagent communication, where in order to solve a task, agents require communicating with another agents.

Original: Ideally, an agent would would like to communicate with other agent, only when its essential to communicate, i.e the agents would like to minimize the communication with another agents.
Revised: Ideally, an agent would would like to communicate with other agent, only when its essential to communicate, i.e the agents would like to minimize the communication with another agents.

Original: Here we show that selectively deciding when to communicate with another agent can result in faster learning.
Revised: Here we show that selectively deciding when to communicate with another agent can result in faster learning.

Original: In order to be more concrete, suppose there are two agents, Alex and Bob, and that Alex receives
Revised: In order to be more concrete, suppose there are two agents, Alex and Bob, and that Alex receives

Original: Bob’s action at time t, then Alex can use Bob’s action to decide what inﬂuence Bob’s action has on its own action distribution.
Revised: Bob’s action at time t, then Alex can use Bob’s action to decide what inﬂuence Bob’s action has on its own action distribution.

Original: Lets say the current state of the Alex and Bob are represented by s a , s r respectively, and the communication channel b/w Alex and Bob is represented by z .
Revised: Lets say the current state of the Alex and Bob are represented by s a , s r respectively, and the communication channel b/w Alex and Bob is represented by z .

Original: Alex and
Revised: Alex and

Original: BOb can decide what action to take based on the distribution p a ( a a | s s ) , p r ( a r | s r ) respectively.
Revised: BOb can decide what action to take based on the distribution p a ( a a | s s ) , p r ( a r | s r ) respectively.

Original: Now, when Alex wants to use the information regarding Bob’s action, then the modiﬁed action distribution (for Alex) becomes p a ( a a | s a , z r ) where z r contains information regarding Bob’s past states and actions, similarly the modiﬁed action distribution (for Bob) becomes p r ( a r | s r , z a ) .
Revised: Now, when Alex wants to use the information regarding Bob’s action, then the modiﬁed action distribution (for Alex) becomes p a ( a a | s a , z r ) where z r contains information regarding Bob’s past states and actions, similarly the modiﬁed action distribution (for Bob) becomes p r ( a r | s r , z a ) .

Original: Now, the idea is that Alex and Bob wants to know about each other actions only when its necessary (i.e the goal is to minimize the communication b/w Alex and Bob) such that on average they only use information corresponding to there own states (which could include past states and past actions.)
Revised: Now, the idea is that Alex and Bob wants to know about each other actions only when its necessary (i.e the goal is to minimize the communication b/w Alex and Bob) such that on average they only use information corresponding to there own states (which could include past states and past actions.)

Original: This would correspond to penalizing the difference between the marginal policy of Alex (“habit policy”) and the conditional policy of Alex (conditioned on Bob’s state and action information) as it tells us how much inﬂuence Bob’s action has on Alex’s action distribution.
Revised: This would correspond to penalizing the difference between the marginal policy of Alex (“habit policy”) and the conditional policy of Alex (conditioned on Bob’s state and action information) as it tells us how much inﬂuence Bob’s action has on Alex’s action distribution.

Original: In mathematical terms, it would correspond to penalizing D KL [ p a ( a a | s a , z r ) | p a ( a a | s a )] .
Revised: In mathematical terms, it would correspond to penalizing D KL [ p a ( a a | s a , z r ) | p a ( a a | s a )] .

Original: In order to show this, we use the same setup as in the paper (Mordatch and Abbeel, 2018).
Revised: In order to show this, we use the same setup as in the paper (Mordatch and Abbeel, 2018).

Original: The agent perform actions and communicate utterances according to policy which is identically instantiated for all the different agents.
Revised: The agent perform actions and communicate utterances according to policy which is identically instantiated for all the different agents.

Original: This policy determines the action, and the communication protocols.
Revised: This policy determines the action, and the communication protocols.

Original: We assume all agents have identical action and observation spaces, and all agents act according to the same policy and receive a shared reward.
Revised: We assume all agents have identical action and observation spaces, and all agents act according to the same policy and receive a shared reward.

Original: We consider the cooperative setting, in which the problem is to ﬁnd a policy that maximizes expected return for all the agents.
Revised: We consider the cooperative setting, in which the problem is to ﬁnd a policy that maximizes expected return for all the agents.

Original: Table 4 shows the training and test rewards as compared to the scenarios when there is no communication b/w different agents, when all the agents can communication with each other, and when there is a KL cost penalizing the KL b/w the conditional policy distribution and marginal policy distribution.
Revised: Table 4 shows the training and test rewards as compared to the scenarios when there is no communication b/w different agents, when all the agents can communication with each other, and when there is a KL cost penalizing the KL b/w the conditional policy distribution and marginal policy distribution.

Original: As evident, agents trained with
Revised: As evident, agents trained with

Original: InfoBot cost achieves the best results.
Revised: InfoBot cost achieves the best results.

Original: D. M ULTIAGENT C OMMUNICATION
Revised: D. M ULTIAGENT C OMMUNICATION

Original: The environment consists of N agents and M landmarks.
Revised: The environment consists of N agents and M landmarks.

Original: Both the agents and landmarks exhibit different characteristics such as different color and shape type.
Revised: Both the agents and landmarks exhibit different characteristics such as different color and shape type.

Original: Different agents can act to move in the environment.
Revised: Different agents can act to move in the environment.

Original: They can also act be effected by the interactions with other agents.
Revised: They can also act be effected by the interactions with other agents.

Original: Asides from taking physical actions, agents communicate with other agents using verbal communication symbols c .
Revised: Asides from taking physical actions, agents communicate with other agents using verbal communication symbols c .

Original: We use the same setup as in the paper (Mordatch and Abbeel, 2018).
Revised: We use the same setup as in the paper (Mordatch and Abbeel, 2018).

Original: E E FFECT ON V ISITATION C OUNT DUE TO EXPLORATION BONUS

Original: [Figure]
Revised: [Figure]

Original:  F C OMPARISON WITH O FF POLICY ALGORITHMS (SAC)
Revised: E E FFECT ON V ISITATION C OUNT DUE TO EXPLORATION BONUS F C OMPARISON WITH O FF POLICY ALGORITHMS (SAC)

Original: In this section, we study in isolation the effect of proposed method as compared to the state of the art off-policy methods like SAC (Haarnoja et al., 2018).
Revised: In this section, we study in isolation the effect of proposed method as compared to the state of the art off-policy methods like SAC (Haarnoja et al., 2018).

Original: For this domain, we again use high value states as an approximation to the goal state following (Goyal et al., 2018).
Revised: For this domain, we again use high value states as an approximation to the goal state following (Goyal et al., 2018).

Original: In order to implement this, we maintain a buffer of 20000 states, and choose the state with the highest value under the current value function, as a proxy to the goal.
Revised: In order to implement this, we maintain a buffer of 20000 states, and choose the state with the highest value under the current value function, as a proxy to the goal.

Original: We compare the proposed method with SAC in sparse reward scenarios.
Revised: We compare the proposed method with SAC in sparse reward scenarios.

Original: We evaluate the proposed algorithm on 4 MuJoCo tasks in OpenAI Gym.
Revised: We evaluate the proposed algorithm on 4 MuJoCo tasks in OpenAI Gym.

Original: (Brockman et al., 2016)
Revised: (Brockman et al., 2016)

Original: The result in Figure 11 shows the performance of the proposed method showing improvement over the baseline on HalfCheetah-v2, Walker2d-v2, Hopper-v2 and Swimmer-v2 sparse reward tasks.
Revised: The result in Figure 11 shows the performance of the proposed method showing improvement over the baseline on HalfCheetah-v2, Walker2d-v2, Hopper-v2 and Swimmer-v2 sparse reward tasks.

Original: We evaluate the performance by evaluating after every 50K steps.
Revised: We evaluate the performance by evaluating after every 50K steps.

Original: We plotted the performance of the proposed method as well as baseline for 500K steps.
Revised: We plotted the performance of the proposed method as well as baseline for 500K steps.

Original: We averaged over 2 random seeds.
Revised: We averaged over 2 random seeds.

Original: [Figure]

Original: G I NFORMATION R EGULARIZATION FOR I NSTRUCTION F OLLOWING
Revised: G I NFORMATION R EGULARIZATION FOR I NSTRUCTION F OLLOWING

Original: Here we use the proposed method in the context of interactive worlds for spatial reasoning where the goal is given by language instruction.
Revised: Here we use the proposed method in the context of interactive worlds for spatial reasoning where the goal is given by language instruction.

Original: The agent is placed in an interactive world [Figure] and agent can take actions to reach the goal speciﬁed by language instruction.
Revised: The agent is placed in an interactive world, and agent can take actions to reach the goal speciﬁed by language instruction.

Original: For ex.
Revised: For ex.

Original: Reach the north-most house, the problem could be challenging because the language instruction is is highly contextdependent.
Revised: Reach the north-most house, the problem could be challenging because the language instruction is is highly context- [Figure] dependent.

Original: Therefore, for better generalization to unseen worlds, the model must jointly reason over the instruction text and environment conﬁguration.
Revised: Therefore, for better generalization to unseen worlds, the model must jointly reason over the instruction text and environment conﬁguration.

Original: We model our task as a Markov Decision Process (MDP), where an agent can take actions to effect the world.
Revised: We model our task as a Markov Decision Process (MDP), where an agent can take actions to effect the world.

Original: The goal to the agent is speciﬁed by the language instruction.
Revised: The goal to the agent is speciﬁed by the language instruction.

Original: The MDP can be represented by the tuple (S,A,G,T,R), where S is the set of all possible states, A is the action set, G is the set of all goal speciﬁcations in natural language, T ( s next | s, a, g ) is the transition distribution, and R(s,g) is the reward function, which is dependent on both the current state as well as goal.
Revised: The MDP can be represented by the tuple (S,A,G,T,R), where S is the set of all possible states, A is the action set, G is the set of all goal speciﬁcations in natural language, T ( s next | s, a, g ) is the transition distribution, and R(s,g) is the reward function, which is dependent on both the current state as well as goal.

Original: A state s ∈ S includes information such as the locations of different entities along with the agent’s own position.
Revised: A state s ∈ S includes information such as the locations of different entities along with the agent’s own position.

Original: Puddle world navigation data In order to study generalization across a wide variety of environmental conditions and linguistic inputs, we follow the same experimental setup as in (Janner et al., 2018).
Revised: Puddle world navigation data In order to study generalization across a wide variety of environmental conditions and linguistic inputs, we follow the same experimental setup as in (Janner et al., 2018) [Figure]

Original: Its basically an extension of the puddle world reinforcement learning benchmark.
Revised: Its basically an extension of the puddle world reinforcement learning benchmark.

Original: States in a grid are ﬁrst ﬁlled with either grass or water cells, such that the grass forms one connected component.
Revised: States in a grid are ﬁrst ﬁlled with either grass or water cells, such that the grass forms one connected component.

Original: We then populate the grass region with six unique objects which appear only once per map (triangle, star, diamond, circle, heart, and spade) and four non-unique objects (rock, tree, horse, and house) which can appear any num- ber of times on a given map.
Revised: We then populate the grass region with six unique objects which appear only once per map (triangle, star, diamond, circle, heart, and spade) and four non-unique objects (rock, tree, horse, and house) which can appear any num- ber of times on a given map.

Original: We compare the proposed method to the UVFA (Schaul et al., 2015).
Revised: We compare the proposed method to the UVFA (Schaul et al., 2015).

Original: We made use of one MLPs and the LSTM to learn low dimensional embeddings of states and goals respectively which are then combined via dot product to give value estimates.
Revised: We made use of one MLPs and the LSTM to learn low dimensional embeddings of states and goals respectively which are then combined via dot product to give value estimates.

Original: Goals are described in text described in text, and hence we use the LSTM over the language instruction.
Revised: Goals are described in text described in text, and hence we use the LSTM over the language instruction.

Original: The state MLP has an identical architecture to that of the UVFA: two hidden layers of dimension 128 and ReLU activations.
Revised: The state MLP has an identical architecture to that of the UVFA: two hidden layers of dimension 128 and ReLU activations.

Original: Text instructions can have both local and global references to objects.
Revised: Text instructions can have both local and global references to objects.

Original: Local references require an understanding of spatial prepositional phrases such as ‘above’, ‘next to’ in order to reach the goal.
Revised: Local references require an understanding of spatial prepositional phrases such as ‘above’, ‘next to’ in order to reach the goal.

Original: This is invariant to the global position of the object references, on the other hand, global references[Table] contains superlatives such as ‘easternmost’ and ‘topmost’, which require reasoning over the entire map.
Revised: This is invariant to the global position of the object references, on the other hand, global references contains superlatives such as ‘easternmost’ and ‘topmost’, which require reasoning over the entire map.

Original: For example, in the case of local reference, one could describe a unique object (e.g. Go to the circle), whereas for global reference might require comparing the positions of all objects of a speciﬁc type (e.g. Go to the northernmost tree).
Revised: For example, in the case of local reference, one could describe a unique object (e.g. Go to the circle), whereas for global reference might require comparing the positions of all objects of a speciﬁc type (e.g. Go to the northernmost tree).

Original: Fig.
Revised: Fig.

Original: 12 shows randomly generated worlds.
Revised: 12 shows randomly generated worlds.

Revised: from https://github.com/JannerM/spatial reasoning

Revised: [Table]

Original: H M INI G RID E NVIRONMENTS FOR O PEN AI G YM
Revised: H M INI G RID E NVIRONMENTS FOR O PEN AI G YM

Original: The FindObj and MultiRoom environments used for this research are part of MiniGrid, which is an open source gridworld package 4 .
Revised: The FindObj and MultiRoom environments used for this research are part of MiniGrid, which is an open source gridworld package.

Original: This package includes a family reinforcement learning environments compatible with the OpenAI Gym framework.
Revised: This package includes a family reinforcement learning environments compatible with the OpenAI Gym framework.

Original: Many of these environments are parameterizable so that the difﬁculty of tasks can be adjusted (eg: the size of rooms is often adjustable).
Revised: Many of these environments are parameterizable so that the difﬁculty of tasks can be adjusted (eg: the size of rooms is often adjustable).

Original: H. T HE W ORLD
Revised: H. T HE W ORLD

Original: In MiniGrid, the world is a grid of size NxN. Each tile in the grid contains exactly zero or one object.
Revised: In MiniGrid, the world is a grid of size NxN. Each tile in the grid contains exactly zero or one object.

Original: The possible object types are wall, door, key, ball, box and goal.
Revised: The possible object types are wall, door, key, ball, box and goal.

Original: Each object has an associated discrete color, which can be one of red, green, blue, purple, yellow and grey.
Revised: Each object has an associated discrete color, which can be one of red, green, blue, purple, yellow and grey.

Original: By default, walls are always grey and goal squares are always green.
Revised: By default, walls are always grey and goal squares are always green.

Original: H. R EWARD F UNCTION
Revised: H. R EWARD F UNCTION

Original: Rewards are sparse for all MiniGrid environments.
Revised: Rewards are sparse for all MiniGrid environments.

Original: In the MultiRoom environment, episodes are terminated with a positive reward when the agent reaches the green goal square.
Revised: In the MultiRoom environment, episodes are terminated with a positive reward when the agent reaches the green goal square.

Original: Otherwise, episodes are terminated with zero reward when a time step limit is reached.
Revised: Otherwise, episodes are terminated with zero reward when a time step limit is reached.

Original: In the FindObj environment, the agent receives a positive reward if it reaches the object to be found, otherwise zero reward if the time step limit is reached.
Revised: In the FindObj environment, the agent receives a positive reward if it reaches the object to be found, otherwise zero reward if the time step limit is reached.

Original: The formula for calculating positive sparse rewards is 1 − 0 .

Original: 9 ∗ ( step _ count/max _ steps ) for FindObj env.

Original: That is, rewards are always between zero and one, and the quicker the agent can successfully complete an episode, the closer to 1 the reward will be.

Original: The max _ steps parameter is different for each environment, and varies depending on the size of each environment, with larger environments having a higher time step limit.

Original: In MultiRoom env, agent only gets reward when it reaches the goal.

Original: H. A CTION S PACE
Revised: H. A CTION S PACE

Original: There are seven actions in MiniGrid: turn left, turn right, move forward, pick up an object, drop an object, toggle and done.
Revised: There are seven actions in MiniGrid: turn left, turn right, move forward, pick up an object, drop an object, toggle and done.

Original: For the purpose of this paper, the pick up, drop and done actions are irrelevant.
Revised: For the purpose of this paper, the pick up, drop and done actions are irrelevant.

Original: The agent can use the turn left and turn right action to rotate and face one of 4 possible directions (north, south, east, west).
Revised: The agent can use the turn left and turn right action to rotate and face one of 4 possible directions (north, south, east, west).

Original: The move forward action makes the agent move from its current tile onto the tile in the direction it is currently facing, provided there is nothing on that tile, or that the tile contains an open door.
Revised: The move forward action makes the agent move from its current tile onto the tile in the direction it is currently facing, provided there is nothing on that tile, or that the tile contains an open door.

Original: The agent can open doors if they are right in front of it by using the toggle action.
Revised: The agent can open doors if they are right in front of it by using the toggle action.

Original: H. O BSERVATION S PACE
Revised: H. O BSERVATION S PACE

Original: Observations in MiniGrid are partial and egocentric.
Revised: Observations in MiniGrid are partial and egocentric.

Original: By default, the agent sees a square of 7x7 tiles in the direction it is facing.
Revised: By default, the agent sees a square of 7x7 tiles in the direction it is facing.

Original: These include the tile the agent is standing on.
Revised: These include the tile the agent is standing on.

Original: The agent cannot see through walls or closed doors.
Revised: The agent cannot see through walls or closed doors.

Original: The observations are provided as a tensor of shape 7x7x3.
Revised: The observations are provided as a tensor of shape 7x7x3.

Original: However, note that these are not RGB images.
Revised: However, note that these are not RGB images.

Original: Each tile is encoded using 3 integer values: one describing the  type of object contained in the cell, one describing its color, and a ﬂag indicating whether doors are open or closed.
Revised: Each tile is encoded using 3 integer values: one describing the https://github.com/maximecb/gym-minigrid type of object contained in the cell, one describing its color, and a ﬂag indicating whether doors are open or closed.

Original: This compact encoding was chosen for space efﬁciency and to enable faster training.
Revised: This compact encoding was chosen for space efﬁciency and to enable faster training.

Original: The fully observable RGB image view of the environments shown in this paper is provided for human viewing.
Revised: The fully observable RGB image view of the environments shown in this paper is provided for human viewing.

Original: H. L EVEL G ENERATION
Revised: H. L EVEL G ENERATION

Original: The level generation in this task works as follows: (1) Generate the layout of the map (X number of rooms with different sizes (at most size Y) and green goal) (2) Add the agent to the map at a random location in the ﬁrst room.
Revised: The level generation in this task works as follows: (1) Generate the layout of the map (X number of rooms with different sizes (at most size Y) and green goal) (2) Add the agent to the map at a random location in the ﬁrst room.

Original: (3) Add the goal at a random location in the last room.
Revised: (3) Add the goal at a random location in the last room.

Original: MultiRoomN X S Y -
Revised: MultiRoomN X S Y -

Original: In this task, the agent gets an egocentric view of its surroundings, consisting of 3 × 3 pixels.
Revised: In this task, the agent gets an egocentric view of its surroundings, consisting of 3 × 3 pixels.

Original: A neural network parameterized as MLP is used to process the visual observation.
Revised: A neural network parameterized as MLP is used to process the visual observation.

Original: FindObjSY - In this task, the agent’s egocentric observation consists of 7 × 7 pixels.
Revised: FindObjSY - In this task, the agent’s egocentric observation consists of 7 × 7 pixels.

Original: We use
Revised: We use

Original: Convolutional Neural Networks to encode the visual observations.
Revised: Convolutional Neural Networks to encode the visual observations.

Original: I M INI P AC M AN E NV
Revised: I M INI P AC M AN E NV

Original: Fuedal Networks primarily has four components: (1) Transition policy gradient, (2) Directional cosine similarity rewards, (3) Goals speciﬁed with respect to a learned representation, and (4) dilated
Revised: Fuedal Networks primarily has four components: (1) Transition policy gradient, (2) Directional cosine similarity rewards, (3) Goals speciﬁed with respect to a learned representation, and (4) dilated

Original: RNN.
Revised: RNN.

Original: For our work we use normal LSTMs, and hence we do not include design choice (4).
Revised: For our work we use normal LSTMs, and hence we do not include design choice (4).