InfoBot: Transfer and Exploration via the Information Bottleneck

A central challenge in reinforcement learning is discovering effective policies for tasks where rewards are sparsely distributed . We postulate that in the absence of useful reward signals , an effective exploration strategy should seek out decision states . These states lie at critical junctions in the state space from where the agent can transition to new , potentially unexplored regions . We propose to learn about decision states from prior experience . 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 . We find that this simple mechanism effectively identifies decision states , even in partially observed settings . In effect , the model learns the sensory cues that correlate with potential subgoals . 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 . 1 INTRODUCTION . Deep reinforcement learning has enjoyed many recent successes in domains where large amounts of training time and a dense reward function are provided . However , learning to quickly perform well in environments with sparse rewards remains a major challenge . Providing agents with useful signals to pursue in lieu of environmental reward becomes crucial in these scenarios . In this work , we propose to incentivize agents to learn about and exploit multi-goal task structure in order to efficiently explore in new environments . We do so by first training agents to develop useful habits as well as the knowledge of when to break them , and then using that knowledge to efficiently probe new environments for reward . We focus on multi-goal environments and goal-conditioned policies ( Foster and Dayan , 2002 ; Schaul et al. , 2015 ; Plappert et al. , 2018 ) . In these scenarios , a goal G is sampled from p ( G ) and the beginning of each episode and provided to the agent . The goal G provides the agent with information about the environment ’ s reward structure for that episode . For example , in spatial navigation tasks , G might be the location or direction of a rewarding state . We denote the agent ’ s policy πθ ( A | S , G ) , where S is the agent ’ s state , A the agent ’ s action , and θ the policy parameters . We incentivize agents to learn task structure by training policies that perform well under a variety of goals , while not overfitting to any individual goal . We achieve this by training agents that , in addition to maximizing reward , minimize the policy dependence on the individual goal , quantified by the conditional mutual information I ( A ; G | S ) . 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 classifiers are trained to achieve high accuracy while simultaneously encoding as little information about the input as possible . This form of “ information dropout ” has been shown to promote generalization performance ( Achille and Soatto , 2016 ; Alemi et al. , 2017 ) . We show that minimizing goal information promotes generalization in an RL setting as well . Our proposed model is referred as InfoBot ( inspired from the Information Bottleneck framework ) . 1 Mila , University of Montreal,2 Mila , McGill University , 3 Deepmind , 4 Google Brain , 5 University College London , 6 University of California , Berkeley . Corresponding author : anirudhgoyal9119 @ gmail.com 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 ) . To see this , note that our regularizer can also be written as I ( A ; G | S ) = Eπθ [ DKL [ πθ ( 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 πθ , DKL is the Kuhlback-Leibler divergence , and π0 ( A | S ) = ∑ g p ( g ) πθ ( A | S , g ) is a “ default ” policy with the goal marginalized out . 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 . Thus , our approach encourages the agent to learn useful behaviours and to follow those behaviours closely , except where diverting from doing so leads to significantly higher reward . Humans too demonstrate an affinity 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 ) . 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. ” Our approach to exploration then involves encouraging the agent to seek out these decision states in new environments . Decision states are natural subgoals for efficient exploration because they are boundaries between achieving different goals ( van Dijk and Polani , 2011 ) . 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 . We encourage the visitation of decision states by first training an agent with an information regularizer to recognize decision states . We then freeze the agent ’ s policy , and use DKL [ πθ ( A | S , G ) | π0 ( A | S ) ] as an exploration bonus for training a new policy . Crucially , this approach to exploration is tuned to the family of tasks the agent is trained on , and we show that it promotes efficient exploration than other task-agnostic approaches to exploration ( Houthooft et al. , 2016 ; Pathak et al. , 2017b ) . S G Z A decoder encoder policy ⇡✓ ( A | S , G ) < latexit sha1_base64= '' vvmsRn44bLbpO2HtyQPL9QhgcNY= '' > AAACDnicbVDJSgNBEO2JW4xb1KOX1hCIIGFGBHOMeNBjRLNAJoSeTk3SpGehu0YIQ77Ai7/ixYMiXj1782/sLAeNPih4vFdFVT0vlkKjbX9ZmaXlldW17HpuY3Nreye/u9fQUaI41HkkI9XymAYpQqijQAmtWAELPAlNb3g58Zv3oLSIwjscxdAJWD8UvuAMjdTNF91YdFMXB4Bs7B66EnwsXVA3ED16e3LlKtEf4HE3X7DL9hT0L3HmpEDmqHXzn24v4kkAIXLJtG47doydlCkUXMI45yYaYsaHrA9tQ0MWgO6k03fGtGiUHvUjZSpEOlV/TqQs0HoUeKYzYDjQi95E/M9rJ+hXOqkI4wQh5LNFfiIpRnSSDe0JBRzlyBDGlTC3Uj5ginE0CeZMCM7iy39J47Ts2GXn5qxQrczjyJIDckRKxCHnpEquSY3UCScP5Im8kFfr0Xq23qz3WWvGms/sk1+wPr4ByeGbPQ== < /latexit > < latexit sha1_base64= '' vvmsRn44bLbpO2HtyQPL9QhgcNY= '' > AAACDnicbVDJSgNBEO2JW4xb1KOX1hCIIGFGBHOMeNBjRLNAJoSeTk3SpGehu0YIQ77Ai7/ixYMiXj1782/sLAeNPih4vFdFVT0vlkKjbX9ZmaXlldW17HpuY3Nreye/u9fQUaI41HkkI9XymAYpQqijQAmtWAELPAlNb3g58Zv3oLSIwjscxdAJWD8UvuAMjdTNF91YdFMXB4Bs7B66EnwsXVA3ED16e3LlKtEf4HE3X7DL9hT0L3HmpEDmqHXzn24v4kkAIXLJtG47doydlCkUXMI45yYaYsaHrA9tQ0MWgO6k03fGtGiUHvUjZSpEOlV/TqQs0HoUeKYzYDjQi95E/M9rJ+hXOqkI4wQh5LNFfiIpRnSSDe0JBRzlyBDGlTC3Uj5ginE0CeZMCM7iy39J47Ts2GXn5qxQrczjyJIDckRKxCHnpEquSY3UCScP5Im8kFfr0Xq23qz3WWvGms/sk1+wPr4ByeGbPQ== < /latexit > < latexit sha1_base64= '' vvmsRn44bLbpO2HtyQPL9QhgcNY= '' > AAACDnicbVDJSgNBEO2JW4xb1KOX1hCIIGFGBHOMeNBjRLNAJoSeTk3SpGehu0YIQ77Ai7/ixYMiXj1782/sLAeNPih4vFdFVT0vlkKjbX9ZmaXlldW17HpuY3Nreye/u9fQUaI41HkkI9XymAYpQqijQAmtWAELPAlNb3g58Zv3oLSIwjscxdAJWD8UvuAMjdTNF91YdFMXB4Bs7B66EnwsXVA3ED16e3LlKtEf4HE3X7DL9hT0L3HmpEDmqHXzn24v4kkAIXLJtG47doydlCkUXMI45yYaYsaHrA9tQ0MWgO6k03fGtGiUHvUjZSpEOlV/TqQs0HoUeKYzYDjQi95E/M9rJ+hXOqkI4wQh5LNFfiIpRnSSDe0JBRzlyBDGlTC3Uj5ginE0CeZMCM7iy39J47Ts2GXn5qxQrczjyJIDckRKxCHnpEquSY3UCScP5Im8kFfr0Xq23qz3WWvGms/sk1+wPr4ByeGbPQ== < /latexit > < latexit sha1_base64= '' vvmsRn44bLbpO2HtyQPL9QhgcNY= '' > AAACDnicbVDJSgNBEO2JW4xb1KOX1hCIIGFGBHOMeNBjRLNAJoSeTk3SpGehu0YIQ77Ai7/ixYMiXj1782/sLAeNPih4vFdFVT0vlkKjbX9ZmaXlldW17HpuY3Nreye/u9fQUaI41HkkI9XymAYpQqijQAmtWAELPAlNb3g58Zv3oLSIwjscxdAJWD8UvuAMjdTNF91YdFMXB4Bs7B66EnwsXVA3ED16e3LlKtEf4HE3X7DL9hT0L3HmpEDmqHXzn24v4kkAIXLJtG47doydlCkUXMI45yYaYsaHrA9tQ0MWgO6k03fGtGiUHvUjZSpEOlV/TqQs0HoUeKYzYDjQi95E/M9rJ+hXOqkI4wQh5LNFfiIpRnSSDe0JBRzlyBDGlTC3Uj5ginE0CeZMCM7iy39J47Ts2GXn5qxQrczjyJIDckRKxCHnpEquSY3UCScP5Im8kFfr0Xq23qz3WWvGms/sk1+wPr4ByeGbPQ== < /latexit > penc ( Z | S , G ) < latexit sha1_base64= '' bZhBgZvFGMw2E2nN/u8tNTfEbf0= '' > AAACDXicdVBNS1tBFJ3nV2PqR9Slm9EoKMjjjU1MshNctEtFE4N5Icyb3JcMzvtg5j4xPPIH3PSvdNNFpXTbfXf+Gycxgi32wMDhnHu5c06QKmnQ856cufmFxaUPheXix5XVtfXSxmbLJJkW0BSJSnQ74AaUjKGJEhW0Uw08ChRcB7dnE//6DrSRSXyFoxS6ER/EMpSCo5V6pb205yPcYw6xGPs7voIQD278SPbp5dFnX8vBEA97pbLnHrMKqzWo51ar7FO9YknjpOLVGWWuN0WZzHDeK/3x+4nIIohRKG5Mh3kpdnOuUQoF46KfGUi5uOUD6Fga8whMN5+mGdN9q/RpmGj7YqRT9e1GziNjRlFgJyOOQ/OvNxHf8zoZhvVuLuM0w0na6aEwUxQTOqmG9qUGgWpkCRda2r9SMeSaC7QFFm0Jr0np/0nr2GWeyy4q5dP6rI4C2Sa75IAwUiOn5As5J00iyAP5Rn6QR+er89356fx6GZ1zZjtb5C84v58BNqKbmg== < /latexit > < latexit sha1_base64= '' bZhBgZvFGMw2E2nN/u8tNTfEbf0= '' > AAACDXicdVBNS1tBFJ3nV2PqR9Slm9EoKMjjjU1MshNctEtFE4N5Icyb3JcMzvtg5j4xPPIH3PSvdNNFpXTbfXf+Gycxgi32wMDhnHu5c06QKmnQ856cufmFxaUPheXix5XVtfXSxmbLJJkW0BSJSnQ74AaUjKGJEhW0Uw08ChRcB7dnE//6DrSRSXyFoxS6ER/EMpSCo5V6pb205yPcYw6xGPs7voIQD278SPbp5dFnX8vBEA97pbLnHrMKqzWo51ar7FO9YknjpOLVGWWuN0WZzHDeK/3x+4nIIohRKG5Mh3kpdnOuUQoF46KfGUi5uOUD6Fga8whMN5+mGdN9q/RpmGj7YqRT9e1GziNjRlFgJyOOQ/OvNxHf8zoZhvVuLuM0w0na6aEwUxQTOqmG9qUGgWpkCRda2r9SMeSaC7QFFm0Jr0np/0nr2GWeyy4q5dP6rI4C2Sa75IAwUiOn5As5J00iyAP5Rn6QR+er89356fx6GZ1zZjtb5C84v58BNqKbmg== < /latexit > < latexit sha1_base64= '' bZhBgZvFGMw2E2nN/u8tNTfEbf0= '' > AAACDXicdVBNS1tBFJ3nV2PqR9Slm9EoKMjjjU1MshNctEtFE4N5Icyb3JcMzvtg5j4xPPIH3PSvdNNFpXTbfXf+Gycxgi32wMDhnHu5c06QKmnQ856cufmFxaUPheXix5XVtfXSxmbLJJkW0BSJSnQ74AaUjKGJEhW0Uw08ChRcB7dnE//6DrSRSXyFoxS6ER/EMpSCo5V6pb205yPcYw6xGPs7voIQD278SPbp5dFnX8vBEA97pbLnHrMKqzWo51ar7FO9YknjpOLVGWWuN0WZzHDeK/3x+4nIIohRKG5Mh3kpdnOuUQoF46KfGUi5uOUD6Fga8whMN5+mGdN9q/RpmGj7YqRT9e1GziNjRlFgJyOOQ/OvNxHf8zoZhvVuLuM0w0na6aEwUxQTOqmG9qUGgWpkCRda2r9SMeSaC7QFFm0Jr0np/0nr2GWeyy4q5dP6rI4C2Sa75IAwUiOn5As5J00iyAP5Rn6QR+er89356fx6GZ1zZjtb5C84v58BNqKbmg== < /latexit > < latexit sha1_base64= '' bZhBgZvFGMw2E2nN/u8tNTfEbf0= '' > AAACDXicdVBNS1tBFJ3nV2PqR9Slm9EoKMjjjU1MshNctEtFE4N5Icyb3JcMzvtg5j4xPPIH3PSvdNNFpXTbfXf+Gycxgi32wMDhnHu5c06QKmnQ856cufmFxaUPheXix5XVtfXSxmbLJJkW0BSJSnQ74AaUjKGJEhW0Uw08ChRcB7dnE//6DrSRSXyFoxS6ER/EMpSCo5V6pb205yPcYw6xGPs7voIQD278SPbp5dFnX8vBEA97pbLnHrMKqzWo51ar7FO9YknjpOLVGWWuN0WZzHDeK/3x+4nIIohRKG5Mh3kpdnOuUQoF46KfGUi5uOUD6Fga8whMN5+mGdN9q/RpmGj7YqRT9e1GziNjRlFgJyOOQ/OvNxHf8zoZhvVuLuM0w0na6aEwUxQTOqmG9qUGgWpkCRda2r9SMeSaC7QFFm0Jr0np/0nr2GWeyy4q5dP6rI4C2Sa75IAwUiOn5As5J00iyAP5Rn6QR+er89356fx6GZ1zZjtb5C84v58BNqKbmg== < /latexit > pdec ( A | S , Z ) < latexit sha1_base64= '' Iq/stS8ZR9f6WF3RymXwlPi7Qrs= '' > AAACDXicdVBNSyNBEO3R9WOzfkT3uJdeswsKMvQY8+FN8eJR0ahsJoSenpqkseeD7hoxDPkDXvavePGgiFfv3vw3dmIW3MV9UPB4r4qqekGmpEHGXpyp6U8zs3Pzn0tfFhaXlssrq6cmzbWAlkhVqs8DbkDJBFooUcF5poHHgYKz4GJ/5J9dgjYyTU5wkEEn5r1ERlJwtFK3/CPr+ghXWIQghv53X0GE63t+LEN6vPnL17LXx41uucLcHVb3ajXK3Gq1XquPyFajyhpN6rlsjAqZ4LBbfvbDVOQxJCgUN6btsQw7BdcohYJhyc8NZFxc8B60LU14DKZTjL8Z0p9WCWmUalsJ0rH6fqLgsTGDOLCdMce++dcbiR957RyjZqeQSZYjJOJtUZQriikdRUNDqUGgGljChZb2Vir6XHOBNsCSDeHPp/T/5HTL9ZjrHW1XdpuTOObJN7JG1olHGmSXHJBD0iKCXJMbckfund/OrfPgPL61TjmTma/kLzhPryDVm40= < /latexit > < latexit sha1_base64= '' Iq/stS8ZR9f6WF3RymXwlPi7Qrs= '' > AAACDXicdVBNSyNBEO3R9WOzfkT3uJdeswsKMvQY8+FN8eJR0ahsJoSenpqkseeD7hoxDPkDXvavePGgiFfv3vw3dmIW3MV9UPB4r4qqekGmpEHGXpyp6U8zs3Pzn0tfFhaXlssrq6cmzbWAlkhVqs8DbkDJBFooUcF5poHHgYKz4GJ/5J9dgjYyTU5wkEEn5r1ERlJwtFK3/CPr+ghXWIQghv53X0GE63t+LEN6vPnL17LXx41uucLcHVb3ajXK3Gq1XquPyFajyhpN6rlsjAqZ4LBbfvbDVOQxJCgUN6btsQw7BdcohYJhyc8NZFxc8B60LU14DKZTjL8Z0p9WCWmUalsJ0rH6fqLgsTGDOLCdMce++dcbiR957RyjZqeQSZYjJOJtUZQriikdRUNDqUGgGljChZb2Vir6XHOBNsCSDeHPp/T/5HTL9ZjrHW1XdpuTOObJN7JG1olHGmSXHJBD0iKCXJMbckfund/OrfPgPL61TjmTma/kLzhPryDVm40= < /latexit > < latexit sha1_base64= '' Iq/stS8ZR9f6WF3RymXwlPi7Qrs= '' > AAACDXicdVBNSyNBEO3R9WOzfkT3uJdeswsKMvQY8+FN8eJR0ahsJoSenpqkseeD7hoxDPkDXvavePGgiFfv3vw3dmIW3MV9UPB4r4qqekGmpEHGXpyp6U8zs3Pzn0tfFhaXlssrq6cmzbWAlkhVqs8DbkDJBFooUcF5poHHgYKz4GJ/5J9dgjYyTU5wkEEn5r1ERlJwtFK3/CPr+ghXWIQghv53X0GE63t+LEN6vPnL17LXx41uucLcHVb3ajXK3Gq1XquPyFajyhpN6rlsjAqZ4LBbfvbDVOQxJCgUN6btsQw7BdcohYJhyc8NZFxc8B60LU14DKZTjL8Z0p9WCWmUalsJ0rH6fqLgsTGDOLCdMce++dcbiR957RyjZqeQSZYjJOJtUZQriikdRUNDqUGgGljChZb2Vir6XHOBNsCSDeHPp/T/5HTL9ZjrHW1XdpuTOObJN7JG1olHGmSXHJBD0iKCXJMbckfund/OrfPgPL61TjmTma/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 : 2 OUR APPROACH . Our objective is to train an agent on one set of tasks ( environments ) T ∼ ptrain ( T ) , but to have the agent perform well on another different , but related , set of tasks T ∼ ptest ( T ) . We propose to maximize the following objective in the training environments : J ( θ ) ≡ Eπθ [ r ] − βI ( A ; G | S ) = Eπθ [ r − βDKL [ πθ ( A | S , G ) | π0 ( A | S ) ] ] , ( 1 ) where Eπθ denotes an expectation over trajectories generated by the agent ’ s policy , β > 0 is a tradeoff parameter , DKL is the Kullback–Leibler divergence , and π0 ( A | S ) ≡ ∑ g p ( g ) πθ ( A | S , g ) is a “ default ” policy with the agent ’ s goal marginalized out . 2.1 TRACTABLE BOUNDS ON INFORMATION . We parameterize the policy πθ ( A | S , G ) using an encoder penc ( Z | S , G ) and a decoder pdec ( A | S , Z ) such that πθ ( A | S , G ) = ∑ z penc ( z | S , G ) pdec ( A | S , z ) ( see figure 1 ) . 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 . The decoder takes this encoded goal information and the current state and produces a distribution over actions A . We suppress the dependence of penc and pdec on θ , but θ in the union of their parameters . 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 penc also minimizes I ( A ; G | S ) . Thus , we instead maximize this lower bound on J ( θ ) : J ( θ ) ≥ Eπθ [ r ] − βI ( Z ; G | S ) = Eπθ [ r − βDKL [ penc ( Z | S , G ) | p ( Z | S ) ] ] , ( 2 ) where p ( Z | S ) = ∑g p ( g ) penc ( Z | S , g ) is the marginalized encoding . 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 . To avoid this marginalization , we replace p ( Z | S ) with a variational approximation q ( Z | S ) ( Kingma and Welling , 2014 ; Alemi et al. , 2017 ; Houthooft et al. , 2016 ; Strouse et al. , 2018 ) . This again provides a lower bound on J ( θ ) since : I ( Z ; G | S ) = ∑ z , s , g p ( z , s , g ) log p ( z | g , s ) p ( z | s ) = ∑ z , s , g p ( z , s , g ) log p ( z | g , s ) − ∑ z , s p ( s ) p ( z | s ) log p ( z | s ) ≥ ∑ z , s , g p ( z , s , g ) log p ( z | g , s ) − ∑ z , s p ( s ) p ( z | s ) log q ( z | s ) , ( 3 ) where the inequality in the last line , in which we replace p ( z | s ) with q ( z | s ) , follows from that DKL [ p ( Z | s ) | q ( Z | s ) ] ≥ 0⇒ ∑ z p ( z | s ) log p ( z | s ) ≥ ∑ z p ( z | s ) log q ( z | s ) . Thus , we arrive at the lower bound J̃ ( θ ) that we maximize in practice : J ( θ ) ≥ J̃ ( θ ) ≡ Eπθ [ r − βDKL [ penc ( Z | S , G ) | q ( Z | S ) ] ] . ( 4 ) In the experiments below , we fix q ( Z | S ) to be unit Gaussian , however it could also be learned , in which case its parameters should be included in θ . Although our approach is compatible with any RL method , we maximize J̃ ( θ ) on-policy from sampled trajectories using a score function estimator ( Williams , 1992 ; Sutton et al. , 1999a ) . As derived by Strouse et al . ( 2018 ) , the resulting update at time step t , which we denote∇θJ̃ ( t ) , is : ∇θJ̃ ( t ) = R̃t log ( πθ ( at | st , gt ) ) − β∇θDKL [ penc ( Z | st , gt ) | q ( Z | st ) ] , ( 5 ) where R̃t ≡ ∑T u=t γ u−tr̃u is a modified return , r̃t ≡ rt + βDKL [ penc ( Z | st , g ) | q ( Z | st ) ] is a modified reward , T is the length of the current episode , and at , st , and gt are the action , state , and goal at time t , respectively . The first term in the gradient comes from applying the REINFORCE update to the modified 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 . 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 . Note that while we take a Monte Carlo policy gradient , or REINFORCE , approach here , our regularizer is compatible with any RL algorithm . In practice , we estimate the marginalization over Z using a single sample throughout our experiments .

Summary: This paper proposes the concept of decision state, which is the state where decision is made “more” dependent to a particular goal. The authors propose a KL divergence regularization to learn the structure of the tasks, and then use this information to encourage the policy to visit the decision states. The method is tested on several different experiment setups.