Title: DYNAMICAL SYSTEM EMBEDDING FOR EFFICIENT INTRINSICALLY MOTIVATED AGENTS Anonymous authors Paper under double-blind review

Abstract: Mutual Information between agent Actions and environment States (MIAS) quantifies the influence of agent on its environment. Recently, it was found that the maximization of MIAS can be used as an intrinsic motivation for artificial agents. In literature, the term empowerment is used to represent the maximum of MIAS at a certain state. While empowerment has been shown to solve a broad range of reinforcement learning problems, its calculation in arbitrary dynamics is a challenging problem because it relies on the estimation of mutual information. Existing approaches, which rely on sampling, are limited to low dimensional spaces, because high-confidence distribution-free lower bounds for mutual information require exponential number of samples. In this work, we develop a novel approach for the estimation of empowerment in unknown dynamics from visual observation only, without the need to sample for MIAS. The core idea is to represent the relation between action sequences and future states using a stochastic dynamic model in latent space with a specific form. This allows us to efficiently compute empowerment with the "Water-Filling" algorithm from information theory. We construct this embedding with deep neural networks trained on a sophisticated objective function. Our experimental results show that the designed embedding preserves information-theoretic properties of the original dynamics.

1. INTRODUCTION:

Deep Reinforcement Learning (Deep RL) provides a solid framework for learning an optimal policy given a reward function, which is provided to the agent by an external expert with specific domain knowledge. This dependency on the expert domain knowledge may restrict the applicability of Deep RL techniques because, in many domains, it is hard to define an ultimate reward function.

On the other hand, intrinsically-motivated artificial agents do not require external domain knowledge but rather get a reward from interacting with the environment. This motivates the study of intrinsically-motivated AI in general, and to develop efficient intrinsically-motivated methods in particular as an alternative and/or complementary approach to the standard reinforcement learning setting.

In recent works, different intrinsic motivation and unsupervised approaches were introduced (Klyubin et al., 2005; Klyubin et al., 2005; Wissner-Gross & Freer, 2013; Salge et al., 2013b; Pathak et al., 2017; Warde-Farley et al., 2018) . Among others, the empowerment method reviewed in Salge et al. (2014) uses the diversity of future states distinguishably-achievable by agent as an intrinsic reward. By definition, empowerment of the agent at a given state is the channel capacity between the agent's choice of action and its resulting next state. Previously, this information-theoretic approach has proved to solve a broad range of the AI tasks (Salge et al., 2014; 2013b; Tiomkin et al., 2017; Klyubin et al., 2005) . Despite the empirical success, the computational burden of the estimation of empowerment is significant, even for known dynamics and fully observable state spaces, as estimating mutual information between high-dimensional random variables is known to be a hard problem (McAllester & Statos, 2018) . Concretely, high-confidence distribution-free lower bounds for mutual information require exponential number of samples. The difficulty in computation has significantly limited the applicability of empowerment in real-life scenarios, ubiquitous in biology and engineering, where an agent observes the environment through sensors (e.g. vision) and does not know an exact dynamical model of the environment.

In this work, we present a novel approach for efficient estimation of empowerment from visual observations (images) in unknown dynamic environments by learning embedding for the dynamic system. The new approach allows us to avoid computationally expensive sampling in high dimensional state/action spaces required for the estimation of MIAS. Efficient computation of empowerment from images opens new directions in the study of intrinsic motivation in real-life environments. As a generic scheme, we structure the interaction between the agent and the environment in a perception-action cycle, shown in Figure 1 . While our proposed scheme is broadly applicable, we demonstrate its usage in two different domains: intrinsically motivated stabilization of non-linear dynamical systems and intrinsically-safe deep reinforcement learning. First, the experiments help us verify that the embedding correctly captures the dynamics in the original space. Additionally, our approach for the estimation of MIAS in the embedded space provides a qualitatively similar MIAS landscape compared to the corresponding one in the original space, which allows us to construct meaningful intrinsic reward for reinforcement learning. Our study of the inverted pendulum is especially important because this non-linear model is a prototype for upright stabilization in various biological and engineering systems.

The paper is organized as follows: in Section 2 we compare our work with previous relevant works. Especially, table 1 provides a thorough comparison between existing methods for the estimation of empowerment, and differentiates our work from the existing state of the art. In Section 3 we provide the necessary formal background to present our approach. In Section 4 we present the method, with the implementation described in Section 5. In Section 6 we demonstrate the proposed method by numerical simulations, and compare to the analytical solution derived in the previous work (Salge et al., 2013a) . Finally, in Section 7 we conclude, and propose promising continuations of this work.

2. RELATED WORKS:

In this section, we compare our approach to the existing approaches for the estimation of empowerment. The criteria for comparison are: (a) the knowledge of the dynamical model, (b) the ability to take high dimensional inputs, (e.g. images/video) and (c) the need for sampling to estimate MIAS. The criteria of our choice are relevant because (a) the true dynamic model is often inaccessible in real-life environments. (b) our IMPAC assumes generic sensor reading, which can be images. (c) sampling of high-dimensional data is expensive and it is desired to have an alternative approach for the estimation of MIAS.

One of the contributions of this work is an efficient representation scheme, achieved by deep neural networks, which does not require sampling for the estimation of mutual information. This is the first work that makes it possible to estimate empowerment from images with unknown dynamics and without sampling.

As shown at Table 1 , the proposed method makes the least assumptions, (unknown dynamics, visual input, no sampling for the estimation of MIAS) compared to existing methods, which makes it applicable to a broader class of problems. In reinforcement learning, we try to find the optimal policy π that maximizes the expected sum of returns along the trajectory:

max θ E H-1 t=0 r(s t , a t ) | π θ .

3.2. EMPOWERMENT:

In this work, we use the established definition of empowerment as a function of state (Klyubin et al., 2005) . The empowerment at state s t is the channel capacity between the action sequence a T -1 t .

= (a t , a t+1 , . . . , a T -1 ) and the final state s T .

Emp(s t ) = max ω(a k-1 t | st) I(a k-1 t ; s k | s t ) = max ω(a k-1 t | st) H(s k | s t ) -H(s k | a k-1 t , s t ) , where I(s k ; a k-1 t | s t )

is the mutual information functional, and ω(a k-1 t | s t ) is a stochastic policy, described by a probability distribution of action trajectories conditioned on s t .

4. PROPOSED APPROACH:

The key idea of our approach is to represent the relation between action sequences and corresponding final states by a linear function in the latent space and to construct a Gaussian linear channel in this space. Importantly, even though the relation between action sequences and future states is linear, the dependency of future states on initial states is non-linear. This representation of a non-linear dynamical system enables us to compute empowerment by an efficient "water-filling" algorithm in the latent space, as explained in Section 4.4.

4.1. INTERACTION MODEL:

We formulate the interaction between the agent and the environment as an MDP (Section 3.1). Each step taken in the environment can be written as:

EQUATION

where s t is the state at time t, a t is the action at time t, and η t is assume to be Gaussian process noise. The agent does not know the dynamical model f , but rather it obtains through its visual sensors, an image observation ν t of the current state s t . Then, after the agent applies an action, a t , the system moves to the next state, s t+1 . The agent has access to all previous observation and action trajectories,

{τ i } N i=1 , τ i = {ν i t , a i t } T t=1

, where N is the number of trajectories and T is the length of each trajectory.

4.2. EMBEDDED SPACES:

From the collected data, {τ i } N i=1 , we embed tuples in the form of (ν i , a i+k-1 i , ν i+k ) where a i+k-1 i = (a i , a i+1 , • • • , a i+k-1 ). Observations and actions are embedded by deep neural networks to the latent spaces Z ∈ R dz and B ∈ R d b , respectively. Here d z is the dimension of the embedded state, and d b is the dimension of the embedded action sequence. These two spaces are related by a linear equation:

EQUATION

where z t is the embedded representation of the image ν t . A(z t ) is a d z by d b state dependent matrix that relates the embedded action sequence, b t+k-1 t to the future embedded state, z t+k . This model was suggested previously by Salge et al. (2013a) in the original state space with known dynamics.

Our key contribution of this work is the embedding of a generic unknown dynamic system into a latent space that satisfy this model. The architecture for the embedding and the objective functions used to train the networks are provided in Section 5.

4.3. INFORMATION CHANNEL IN EMBEDDED SPACE:

To compute mutual information between embedded action sequences, b t+k-1 t and embedded states, z t+k , we assume that the noise in the system comes from a Gaussian in latent state space: η ∼ N (0 dz×dz , I dz×dz ). As a result, we end up with an Gaussian linear channel in embedded spaces:

EQUATION

This particular formulation of dynamic systems in the latent space allows us to solve for the channel capacity directly and efficiently, as described in below.

4.4. CHANNEL CAPACITY IN LATENT SPACE:

To compute the information capacity of the channel given in Eq. 4.3, we apply the "Water-Filling" algorithm (Cover & Thomas, 2012) , which is a computationally efficient method for the estimation of capacity, C * , of a linear Gaussian channel:

EQUATION

under the constraint on the total power, n i=1 p i = P , where {σ i (z t )} n i=1 are the singular values of the channel matrix, A(z t ), and P is the total power in the latent space.

5. ARCHITECTURE:

The most crucial component of the IMPAC scheme is the embedding from sensor inputs (in our case, images) and action sequences, into latent state and action tensors while remaining consistent with the real environment dynamics. In this section, we organize our discussion around how a single sample: {S t , a t , • • • , a t+k-1 , S t+k } fits in the training of our deep neural networks, as illustrated in Figure 2 . We will explain each individual components of the block diagram and wrap up by discussing the whole picture.

5.1. EMBEDDING OF IMAGE OBSERVATIONS:

The left side and right side of the block diagram in Figure 2 embeds raw observations (images).

In this work, we take images scaled down to 64 × 64 as observations. The image resolution is chosen as a balance of clarity and performance: visual features of the environments are retained while training batches can be fit on a single GPU. Multiple images can be concatenated to capture additional information (e.g. velocity). When encoding multi-channel images into latent vectors, we use an auto-encoder and decoder pair parameterized by convolutional neural networks. Given a training sample {S t , a t , • • • , a t+k-1 , S t+k }, both S t and S t+k are encoded into their respective latent vectors Z t and Z t+k using a same CNN encoder. To ensure that essential state information is conserved after embedding, we decoded the latent tensor Z t back to images using a deconvolutional network for correct reconstruction Ŝt . The objective is to make the images S t and Ŝt similar. The encoder and decoder network details are presented in Appendix B.

5.2. EMBEDDING OF ACTION SEQUENCES:

The mid-lower section of Figure 2 embeds action sequences.

Standard MDPs (Section 3.1) models environment dynamics in the granularity of single steps, which results in an exponentially large action space over long horizons. Another contribution of this work is that we compute information-theoretic quantities across multiple time-steps without exploding the action space. We consider a k-step dynamic model directly by embedding action sequence {a t , a t+1 • • • a t+k-1 } into a single action b t in the latent space (1 step in the latent space corresponds to k steps in the original space). Actions inside the sequence are concatenated into a vector and passed into an auto-encoder to get b t . Again, to ensure that necessary information is retained in the latent action tensor, we decode b t for reconstruction of {a t , a t+1 • • • a t+k-1 }. In our implementation, the encoder and decoder are both fully connected MLPs. The detailed information about each layer can be found in Appendix B.

5.3. DYNAMICAL MODEL IN LATENT SPACE:

The top part of Figure 2 fits a dynamics model linear in action (non-linear in state).

As discussed in Section 4.3, our scheme requires a dynamic model with a particular form. Specifically, given current latent state Z t , the next state Z t+k is linear in the action b t . Mathematically, if latent states have dimension d z and latent action have dimension d b , we want to find a mapping

A : R dz → R dz×d b such that Z t+k ≈ A(Z t ) × b.

The mapping from Z t to an d z × b z matrix is parameterized by a fully connected MLP with d z × d b output neurons. The final layer is then reshaped into a matrix. (See appendix B)

5.4. THE ENTIRE SCHEME:

In the previous 3 sub-sections, we introduce 3 objectives: 1) Find an embedding of observations that retains their information. 2) Find an embedding of action sequences that retain their information. 3) Find a dynamics model linear in the latent action. When these three objectives are trained jointly, the neural networks learn towards a latent space that extract information from the original space and learns a dynamics model that is consistent with the actual environment.

L = αE (S t -Ŝt ) 2 + βE (a t -ât ) 2 + γE (Z t+k -Ẑt+k ) 2

(5)

In practice, regularization on the latent tensors is enforced at training time to encourage a structured latent space. The details about our choice of regularization are discussed in appendix A.

6. EXPERIMENTS:



6.1. INVERTED PENDULUM:

Inverted pendulum is a prototypical system for upright bi-pedal walking, and it often serves to check new AI algorithms. Moreover, inverted pendulum is a baseline for the estimation of empowerment in dynamical systems because its landscape is well-studied in numerous previous studies (Salge et al., 2013b; 2014; 2013a; Karl et al., 2017; Mohamed & Rezende, 2015) . So, we chose inverted pendulum from the Openai Gym environment for testing our new method, and comparing it to the previous methods (Brockman et al., 2016) .

6.1.1. CORRECTNESS OF LATENT DYNAMICS:

Prior to applying the proposed model, given by Eq. 4.3, to the estimation of MIAS in the latent space, we verify the correctness of our embedded dynamic model. We compare reconstruction of latent prediction with ground truth observation. In inverted pendulum, two consecutive images are used to capture the state information (angle and angular velocity). When we supply the initial two images for time t, along with choices of different action sequences of length k, we expect the reconstructions to match the corresponding observations at time t + k. As shown in Figure 3 the model is capable of reconstructing the future images from the current images and action sequences. 

6.1.2. EMPOWERMENT LANDSCAPE AND PENDULUM SWING-UP:

In this section, we compute the empowerment values and use those as intrinsic reward to train an optimal trajectory for swinging up the pendulum from the bottom position to the upright position.

The optimal trajectory is computed by the standard PPO algorithm with reward signal coming solely from empowerment. (Schulman et al., 2017b) . Figure 4 compares the empowerment landscape computed under our approach with a previous result:

As shown in Figure 4 , at the beginning, the pendulum swings with a small amplitude (the dense black dots around θ = π rad, θ = 0 rad s -1 ). Once it accumulated enough energy, the pendulum starts traversing the state space, and arrives at the top position, where it stabilises, (dense black points around θ = 0 rad, θ = 0 rad s -1 ).

Figure 4 : Left: Optimal trajectory super-imposed on our empowerment landscape. The pendulum swings up from bottom (π, 0) to balance at the top (0, 0).Right: Empowerment values calculated with knowledge of the ground truth model (Salge et al., 2013a) . This figure shows that the empowerment landscapes are qualitatively similar, which verifies the validity of our approach.

Results from the pendulum experiment effectively demonstrate the validity of our proposed approach. Not only was our method able to create an empowerment value plot similar to that of a previous analytical method, but also these empowerment values train to balance the pendulum at upright position. The pendulum experiment shows that our latent representation effectively captures the information-theoretic properties of the original environment.

6.2. SAFETY OF RL AGENT:

Another useful application of empowerment is its implication of safety of the artificial agent. A state is intrinsically safe for an agent when the agent has a high diversity of future states, achievable by its actions. This is because in such states, the agent can take effective actions to prevent undesirable futures. In this context, the higher its empowerment value, the safer the agent is. In this experiment, we first check that our calculation of empowerment matches the specific design of the environment. Additionally, we show that empowerment augmented reward function can affect the agent's preference between a shorter but more dangerous path and a longer but safer one.

Environment: a double tunnel environment implemented with to the OpenAI Gym API (Brockman et al., 2016) . Agent (blue) is modeled as a ball of radius 1 inside a 20×20 box. The box is separated by a thick wall (gray) into top and bottom section. Two tunnels connect the top and bottom of the box. The tunnel in middle is narrower but closer to the goal compared to the one on the right.

Control::

In each time step, the agent can move at most 0.5 unit length in each of x and y direction.

If an action takes the agent into the walls, it will be shifted back out to the nearest surface.

Reset criterion: each episode has a maximum length of 200 steps. The environment resets when time runs out or when the agent reaches the goal.

Since the tunnel in the middle is narrower, the agent is relatively less safe there. The effectiveness of the control of the agent is damped in 2 ways:

1. In a particular time step, it's more likely for the agent to bump into the walls. When this happens, the agent is unable to proceed as far as desired.

2. The tunnel is 10 units in length. When the agent is around the middle, it will still be inside the tunnel in the 5-step horizon. Thus, the agent has fewer possible future states.

We trained the agent with PPO algorithm (Schulman et al., 2017a) from OpenAI baselines (Dhariwal et al., 2017) using an empowerment augmented reward function. After a parameter search, we used discount factor γ = 0.95 over a total of 10 6 steps. The reward function that we choose is: R(s, a) = 1 goal + β × Emp(s) The results from the tunnel environment again support our proposed scheme. First, the empowerment landscape matches our design of the environment. Second, high quality empowerment reward successfully alters the behavior of the agent.

7. DISCUSSION:

Intrinsically motivated artificial agents do not rely on external reward and thus, do not need domain knowledge for solving a broad class of AI problem such as stabilization, tracking, etc. In this work, we introduce a new method for efficient estimation of a certain type of information-theoretic intrinsic motivation, known as empowerment. The learnt embedding representation reliably captures not only the dynamic system but also its underlying information-theoretic properties.

With this work, we also introduced the IMPAC framework, which allows for a systematic study of intrinsically motivated agents in broader scenarios, including real-world physical environments like robotics. Future works on physical environment will better demonstrate the usefulness and general applicability of empowerment as an intrinsic motivation. 

:



Figures and Tables:

Figure 1: Intrinsically Motivated Perception and Action Cycle. Left-to-right: (a) the observations with the corresponding action sequences are embedded to the latent space by DNNs, where the final states are constrained to be linear in the action sequences (non-linear in the initial states). (b) Gaussian channel capacity between the embedded actions and the embedded final states is computed efficiently by the 'Water-filling' algorithm. (c) an intrinsic reward is derived, triggering the agent to improve its policy for better reward. (d) when the agent operates under the new policy, additional observations from the environment are perceived by the sensors, closing the loop.

Figure 2: An illustration of data flow around a single training sample.

Figure 3: Reconstruction through the latent dynamics in comparison with the ground truth. Starting from the same upright position, a different action sequence is taken in each row. In all 3 cases, reconstruction matches the actual observation after the actions are taken.

Figure 5: Left sub-figure: Double tunnel environment. The goal is marked in red. The agent in blue. Right sub-figure: Empowerment landscape for the tunnel environment. The values of empowerment reduce at the corner and inside the tunnel where the control of the agent is less effective compared to more open locations.

(a) Observation reconstruction error: L obs = ||S t -St || 2 2 (b) Action sequence reconstruction error: L action = ||a k t -ãk t || 2 2 (c) Prediction error in latent space: L latent = ||Z t+k -Ẑt+k || 2 2 (d) Prediction error in original space: L org = ||S t+k -St+k || 2 2 8. In additional to the loss terms, we add regularization terms to prevent latent vectors from shrinking. This help us get consistent and comparable empowerment values across different trials and even different environments. (a) Regularization of latent state: Reg z = |1 -Regularization of latent action: Reg b = |1 -, we use batched gradient descent on the overall loss function to train all the neural networks in the same loop. L = α obs L obs + α action L action + α latent L latent + α org L org + α reg (Reg z + Reg b ) For both of our experiments, we chose α obs = α org = 100 α action = 10 α latent = α reg = 1 and were able to produce desired empowerment plots. B DETAILS ON NEURAL NETWORK LAYERS B.1 CONVOLUTIONAL NET FOR IMAGE ENCODING (h1) 2D convolution: 4 filters, stride 2, 32 channels, ReLU activation (h2) 2D convolution: 4 filters, stride 2, 64 channels, ReLU activation (h3) 2D convolution: 4 filters, stride 2, 128 channel, ReLU activations (h4) 2D convolution: 4 filters, stride 2, 256 channel, ReLU activations (out) Flatten each sample to a 1D tensor of length 1024 B.2 DECONVOLUTIONAL NET FOR IMAGE RECONSTRUCTION (h1) Fully connected layer with 1024 neurons, ReLU activation (h1') Reshape to 1 × 1 images with 1024 channels (h2) 2D conv-transpose: 5 filters, stride 2, 128 channels, ReLU activation (h3) 2D convolution: 5 filters, stride 2, 64 channels, ReLU activation (h4) 2D convolution: 6 filters, stride 2, 32 channel, ReLU activations (out) 2D convolution: 6 filters, stride 2, C channel B.3 MLP FOR ACTION SEQUENCE ENCODING (h1) Fully connected layer with 512 neurons, ReLU activation (h2) Fully connected layer with 512 neurons, ReLU activation (h3) Fully connected layer with 512 neurons, ReLU activation (out) Fully conected layer with 32 output neurons B.4 MLP FOR ACTION SEQUENCE RECONSTRUCTION (h1) Fully connected layer with 512 neurons, ReLU activation (h2) Fully connected layer with 512 neurons, ReLU activation (out) Fully conected layer with k × d a neurons, tanh activation then scaled to action space B.5 MLP FOR TRANSITION MATRIX A (h1) Fully connected layer with 1024 neurons, ReLU activation (h2) Fully connected layer with 4096 neurons, ReLU activation (h2) Fully connected layer with 8192 neurons, ReLU activation (out) Fully conected layer with d z × d b (1024 × 32) neurons

Comparison between existing methods for the estimation of empowerment.

Table content: <table><tr><td>Method</td><td colspan="3">Unknown dynamics Visual input Sampling</td></tr><tr><td>Salge et al. (2013a)</td><td>no</td><td>no</td><td>no</td></tr><tr><td colspan="2">Mohamed &amp; Rezende (2015) yes</td><td>yes</td><td>yes</td></tr><tr><td>Gregor et al. (2016)</td><td>yes</td><td>yes</td><td>yes</td></tr><tr><td>Karl et al. (2017)</td><td>yes</td><td>no</td><td>yes</td></tr><tr><td>Tiomkin et al. (2017)</td><td>no</td><td>no</td><td>no</td></tr><tr><td>Our work</td><td>yes</td><td>yes</td><td>no</td></tr><tr><td>3 PRELIMINARIES</td><td/><td/><td/></tr><tr><td>3.1 MARKOV DECISION PROCESS</td><td/><td/><td/></tr></table>

A DETAILS OF TRAINING ALGORITHM 1. Collect trajectories using random policy and random reset. 2. Separate out tuples (S t , a t , a t+1 , • • • , a t+k-1 , S t+k ) from trajectories. For simplicity of notation, let a k t represent the concatenated action sequences a t , a t+1 , • • • , a t+k-1 3. Use randomly initialized encoder networks to map the observations and action sequence into latent space. This gives us tuples (Z t , b k t , Z t+k ). 4. Use randomly initialized MLP network to get the corresponding transformation matrices A(Z t ). Ẑt+k ) Note that St+k is reconstructed from latent prediction. 7. Calculate the following loss terms:

Table content: <table><tr><td>5. Calculate the predicted next state in latent space Ẑt+k = A(Z t )b k t</td></tr><tr><td>6. Use randomly initialize decoder networks to reconstruct images and action sequences from</td></tr><tr><td>latent vectors.</td></tr><tr><td>St = Dec(Z t )</td></tr><tr><td>ãk t = Dec(b k t )</td></tr><tr><td>St+k = Dec(</td></tr></table>

Peer review (iclr2020):
This paper proposes to take advantage of a known result on the channel capacity of a linear-Gaussian channel in order to estimate the empowerment and maximize mutual information between policies (action sequences) and final states (given initial states). The idea is to map the raw action sequences and states to a latent space where learning would force that linear property to be appropriate. I like the general idea of the paper (as stated above) along with its objectives but I have several concerns. First, I need to be reassured that we are computing the right quantity. Channel capacity is the maximum mutual information (between inputs and outputs) over the input distribution, whereas I had the impression that empowerment would be this mutual information, and that we want to increase it, but not necessarily reach its maximum over all possible policies: it would usually be one of the terms in an objective function (e.g. here we have reconstruction error, and in practice there would be some task to solve in addition to the exploration reward). One way to see this problem in the given formulation is that the C* objective only depends on the matrix A (which encapsulates the conditional density of z_{t+1} given the trajectory) and it does not depend at all on the distribution of the trajectory itself! This is weird since if we are going to use this as reward the objective is to improve the trajectory. What*s the catch? So either I misunderstand something (which is quite possible) or there is something seriously wrong here. I am assuming that the training objective for the encoder is a sum of the reconstruction error and of C*. But note how this does not give a reward for policies, as such. This is a bit strange if the goal is to construct an exploratory reward! A less radical comment is: have you verified that the linear relationship between z_{t+k} and the b actually holds well? In other words, is the encoder able to map the raw state and actions to a space where the linearity assumption is correct, and thus where equation (3) is satisfied. Figure 3 has something weird, probably one of the two sequences of -1*s should be a sequence of +1*s. Figure 4 is difficult to interpret, the caption should do a better job. The experiment on the safety of the RL agent is weak. I don*t see the longer path as safer, here. And the results are not very impressive, since the agent is only doing what it*s told, i.e., go in areas with more options (more open areas) but there is no reason to believe that this is advantageous, here. Finally, what would make this paper much more appealing is if the whole setup led to learning better high-level representations (how to measure that is another question, but it is a standard kind of question in representation learning papers). Related work: I don*t understand why in the abstract the authors refer to sampling-based methods as requiring exponentially many samples. This is not generally the case for sampling based methods (e.g. think of VAEs). I suppose the authors refer to something in particular but it was not clear to me what. References: in the intro, you might want to refer to the contrastive methods and variational methods to maximize mutual information between representations of the state and representations of the actions, e.g., Thomas et al, 2018, 1802.09484 Kim et al, 2018, arXiv:1810.01176 Warde-Farley et al, 2018, arXiv:1811.11359 Evaluation: for now I suggest a weak reject but I am ready to modify my score if I am convinced that my main concerns were unfounded.