Learning Discrete State Abstractions With Deep Variational Inference ∗
1. Introduction
High-dimensional state spaces are common in deep reinforcement learning (Mnih et al., 2015).
Although states may be as large as images, typically the information required to make good decisions is much smaller.
This motivates the need for state abstraction , the process of encoding states into compressed representations that retain features that inform action and discard uninformative features.
One principled approach to state abstraction is bisimulation in Markov decision processes (MDP) (Dean and Givan, 1997).
Bisimulations formalize the notion of ﬁnding a smaller equivalent abstract MDP that preserves transition and reward information; i.e., it retains relevant decision-making information while reducing the state space size.
We demonstrate this idea in Figure 1, where a grid world with ﬁfteen states is compressed into an MDP with three states.
In this paper, we introduce an approach to ﬁnding approximate MDP bisimulations using the variational information bottleneck (VIB, Tishby et al.
; Alemi et al.
The
VIB framework is typically used to learn representations that predict quantities of interest accurately while ignoring certain aspects of the domain.
VIB methods have previously been applied to state abstraction, but the learned abstraction does not in general take the form of an MDP bisimulation (Abel et al., 2019).
This is problematic, because the abstract MDP can only represent the policies it was trained on, but cannot be used to plan for new tasks.
To resolve this, whereas Abel et al.
(2019) use the abstract states to predict actions from an expert policy, we use abstract states to predict learned Q-values in the VIB objective.
In our setup, a learned encoder maps a state s in the original MDP into a continuous embedding z .
We map the continuous state z onto a discrete abstract state ¯ s by performing inference in a learned probabilistic model.
Our VIB method learns an encoder (i.e. a state abstraction s (cid:55)→ z ) that is predictive of the Q-values that are returned by a deep Q-network, but is regularized using structured prior over the embedding space ( z ).
Concretely, we propose using priors that prefer clusters with Markovian transition structure.
A sequence of embedded states ( z 1 , z 2 , . . . ) is treated as observations from either a Gaussian mixture model
[Figure] [Figure]
(GMM) or an action-conditioned hidden Markov model (HMM), where each embedding z t is emitted from a latent cluster representing abstract state ¯ s t .
In the HMM case, we also learn a cluster transition matrix for each action, serving as the abstract MDP transition model.
The key insight is that abstract states ¯ s group together ground states s (and embeddings z ) with similar Q-values and similar transition properties, thereby forming an approximate
MDP bisimulation.
In addition to the neural encoder, the parameters of our GMM and HMM priors are learned as well.
The learned parameters (cluster means, covariances, and discrete transition matrix between clusters) therefore form our abstract MDP state space and transition function.
When presented with tasks not seen during training, we can use the learned abstract model to plan to solve these tasks without additional learning eﬃciently.
2. Learning bisimulations
We propose a variational method for ﬁnding bisimulations directly from experience.
The end result of our process is an abstract MDP, in which we can eﬃciently plan policies.
Our model consists of three parts:
A deep neural encoder q ( z | s ) that projects states (usually represented as images) onto a low-dimensional continuous latent space (Figure 2 left).
A generative model p ( z, ¯ s ) with a tractable posterior p (¯ s | z ), that encodes the prior belief that the experience was generated by a small discrete Markov process (Figurelower right).
A linear decoder p ( y | z, a ) that predicts state-action values from the continuous encodings (Figure 2 upper right).
We tie the three models together using the deep variational information bottleneck method (Alemi et al., 2017).
2.1. Bisimulation as an information bottleneck
We apply deep IB methods to learn bisimulations as follows.
Let q ( s t , a, y, s t +1 ) be a empirical distribution representing a dataset of transitions from state s t to state s t +1 under an action a , selected by an arbitrary policy π .
y
=
Q π ( s t , a ) denotes the state-action value for the pair ( s t , a ) under π .
We will use the IB method to ﬁnd a compact latent encoding of s t that enables us to predict y while simultaneously matching a prior on the temporal dynamics of the process that generated the data.
Let s = ( s t , s t +1 ) denote a sequential pair of states and z =
( z t , z t +1 ) a corresponding sequential pair of latent states.
The standard IB formulation is:
[Equation]
where use q ( s, a, y, z ) =
q ( s t , a, y, s t +1 )
q ( z | s, a ) as shorthand notation and expand s =
( s t , s t +1 ) and z = ( z t , z t +1 ) in the last identity.
We make two architectural decisions grounded in standard Markov assumptions.
First, we assume that the value y is conditionally independent of z t +1 given z t : q ( y | z t , z t +1 , a ) =
q ( y |
z t , a ).
Second, we assume that z t is conditionally independent of z t +1 , s t +1 , and a given s t and likewise that z t +1 is conditionally independent of z t , s t , and a : q ( z t , z t +1 | s t , s t +1 , a ) =
q ( z t | s t ) q ( z t +1 | s t +1 ).
Together, these two assumptions yield an IB lower bound of the form
[Equation]
L IB presents a trade-oﬀ between encoding enough information of s t in order to predict y (the ﬁrst term of Equation 2) and making the sequence ( z t , z t +1 ) likely under our prior (the second term).
This prior, p ( z t , z t +1 | a ), is a key element of our approach and is discussed in the next section.
Notice that log p ( y | z t , a ) (the ﬁrst term in Equation 2) predicts a state-action value, not a reward.
This provides additional supervision.
Without it, the model tends to collapse all states where the agent does not receive a reward (if the reward function is sparse) into a single abstract state.
2.2. Structured Priors
The denominator of the second term in Equation 2 p ( z t , z t +1 | a ) is the prior.
We use this prior to incorporate an inductive bias, which is that we are observing a discrete Markov process.
We evaluate two priors for this purpose: a prior based on a Gaussian mixture model and one based on an action-conditioned Hidden Markov model.
2.2.1. GMM Prior
We assume K components, each parameterized by a mean µ k and a covariance Σ k
[Equation]
Here ρ k = p ( c t = k ) denotes the probability that z t was generated by component k and
N ( z t | µ k , Σ k ) =
p ( z t | c t = k ) is the Gaussian distribution for the k th component.
In this paper, we constrain Σ k to be diagonal.
For the GMM prior, we set p ( z t , z t +1 | a ) =
p gmm ( z t ) and allow µ k and Σ k to vary; ρ k is uniform and ﬁxed.
This encodes a desire to ﬁnd a latent encoding generated by membership in a ﬁnite set of discrete states (the mixture components).
Each mixture component corresponds to a distinct abstract state.
The weighting function ρ k is the probability that the continuous encoding z t was generated by the k th abstract state.
2.2.2. HMM Prior
To capture the temporal aspect of a Markov process, we can model the prior as an action conditioned hidden Markov model (an HMM).
Here, the “hidden” state is the unobserved discrete abstract state c t used to generate “observations” of the latent state z t .
As in the
GMM, there are K discrete abstract states, each of which generates latent states according to a multivariate Normal distribution with mean µ k and (diagonal) covariance matrix Σ k .
Since we are modelling a Markov process, we include a separate transition matrix T a for each action a where T ak,l denotes the probability of transitioning from an abstract state k to an abstract state l under an action a .
Using this model, the prior becomes:
[Equation]
As with the GMM prior, we allow the parameters of this model ( µ k , Σ k , and T a ) to vary, except for ρ k , which is uniform and ﬁxed.
The transition model T a found during optimization is a discrete conditional probability table that deﬁnes a discrete abstract MDP.
Essentially, this method ﬁnds the parameters of a hidden discrete abstract MDP that ﬁts the observed data over which the loss of Equation 2 is evaluated.
2.3. Deep encoder and end-to-end training
The lower bound L IB (Equation 2) is deﬁned in terms of three distributions that we need to parameterize: the encoder q ( z | s ), the Q-predictor p ( y |
z t , a ) and the prior p ( z t , z t +1 | a ).
The encoder q ( z | s ) is a convolutional network that predicts the mean µ CNN ( s t ) and the diagonal covariance Σ CNN ( s t ) of a multivariate normal distribution.
We used a modiﬁed version of the encoder architecture from (Ha and Schmidhuber, 2018) 1 : ﬁve convolutional layers followed by one fully-connected hidden layers and two fully-connected heads for µ CNN and Σ CNN , respectively.
The Q-predictor p ( y | z t , a ) is a single fully-connected layer (i.e. a linear transform).
We chose this parameterization to impose another constraint on the latent space: the encodings z not only need to form clearly separable clusters to adhere to the prior, but also linearly dependent on their state-action values for each action.
When we train on state-action values for multiple tasks, we predict a vector y instead of a scalar.
Using the reparameterization trick to sample from q ( z | s ), we can compute the gradient of the objective with respect to the encoder weights, Q-predictor weights and the prior parameters.
The prior parameters include the component means and variance, together with the transition function for hidden states in the HMM.
2.4. Planning in the Abstract MDP
A key aspect of our approach is that we can solve new tasks in the original problem domain by solving a compact discrete MDP for new policies.
This is one of the critical motivations for using bisimulations: optimal policies in the abstract MDP induce optimal policies in the original MDP (Givan et al., 2003).
We deﬁne the abstract MDP ¯ M = (cid:104) ¯ S, A, ¯ T , ¯ R, γ (cid:105) with the discrete transition table learned by the HMM prior.
The abstract reward function can encode any reward function in the ground MDP by projecting ground rewards into the abstract space using the encoder.
Now, we can use standard discrete value iteration to ﬁnd new policies.
These policies can be immediately applied in the ground MDP: observations of state in the ground MDP can be projected into the discrete abstract MDP and the new policy can be used to calculate an action.
3. From VIB abstraction to bisimulation
The HMM embedded in our model learns parameters of a compact discrete MDP (Subsection 2.2), but it is not a priori clear that this abstraction is also a bisimulation.
We show that under idealized conditions, every optimal solution to the objective L IB is a bisimulation.
Under assumptions state in Appendix C, the following theorem holds.
Theorem
There exist model parameters θ that reach the global maximum of L IB ( θ ) = − β log K with β > 0 .
The abstraction mapping φ induced by any such model parameters is a bisimulation.
4. Shapes World Experiment
We use a modiﬁed version of the Pucks World from Biza and Platt (2019).
The world is divided into a 4 × 4 grid and objects of various shapes and sizes can be placed or stacked in each cell.
States are represented as simulated 64 × 64 depth images.
The agent can execute a PICK or PLACE action in each of the cells to pick up and place objects.
The goal of abstraction is to recognize that the shapes and sizes of objects do not inﬂuence the PICK and PLACE actions.
We instantiate eight diﬀerent tasks described in Figure 4.
[Table]
We test the ability of the learned abstract models to plan for new goals.
We are able to reach a goal only if it is represented as a distinct abstract state in our model–such abstract states can only exist if the training dataset contains examples of the goal.
Therefore, we can generalize to unseen goals in the sense that our model does not know about these goals during training, but they are represented in the dataset.
During planning for a particular goal, we create a new reward function for the abstract model and assign a reward 1 to all transitions in the dataset that reach that goal.
we run Value Iteration in the abstract model and the found
See Appendix B.3 for more details.
Our model is trained on one or two tasks and we report its ability to plan for every single task (Table 1).
For tasks with a moderate number of abstract states (e.g. 2 objects stacking in a 4 × 4 grid world has 136 abstract states in the coarsest bisimulation), our method can successfully transfer to new tasks of similar complexity without additional training.
For instance, the abstract model learned from two pucks stacking can plan for placing two and three pucks in a row with a 90%+ success rate.
The middle section of Table 1 shows tasks with coarsest abstractions at the limit of what our abstract model can represent.
We can still transfer to similar tasks with a success rate higher than 75%.
The bottom section of Table 1 demonstrates that our algorithm can ﬁnd partial solutions even if the number of abstract states in the coarsest bisimulation exceeds the capacity of the HMM prior.
See Appendix B for additional experiments.
Acknowledgments
Special thanks to Yea-Eun Jang for graphical design in the paper and presentation materials.
The authors also thank members of the GRAIL and Helping Hands groups at Northeastern, as well as anonymous reviewers, for helpful comments on the manuscript.
This work was supported by the Intel Corporation, the 3M Corporation, National Science Foundation (1724257, 1724191, 1763878, 1750649, 1835309), NASA (80NSSC19K1474), startup funds from Northeastern University, the Air Force Research Laboratory (AFRL) and DARPA.
Appendix A. Related Work
Li et al.
(2006) analyzed the topology of the space of abstractions, including bisimulations.
Further analysis by Abel et al.
(2016) proved bounds on the regret of planning an optimal policy in an abstracted MDP compared to the ground MDP.
A work more closely related to ours uses a deep neural net to learn the bisimulation metric between states represented by images (Castro, 2019).
The main diﬀerence between this line of work and ours is that we aim to learn an abstract MDP with discrete states, in which we can plan eﬃciently, whereas bisimulation metric is more commonly used to ﬁnd similar states for the purpose of transfer of policies (Castro and Precup, 2010, 2011) or as an auxiliary loss (Gelada et al., 2019).
Several recent neural-net-based methods use discrete representation for planning.
Serban et al.
(2018) trained a factored transition given a predeﬁned discrete state abstraction.
They focused their empirical evaluation on Natural language processing tasks.
Kurutach et al.
(2018) proposed a Generative adversarial network for learning a forward model with either continuous or discrete latent states.
While they show superior performance on a rope manipulation task with continuous latent states, the discrete state representation learning and planning was only evaluated on a toy 2D navigation task.
Corneil et al.
and van der Pol et al.
(2020) both learn an abstract MDP with discrete states based on ground states represented as images.
Corneil et al.
(2018) used variational inference (Kingma and Welling, 2014) and the Concrete distribution reparameterization trick (Maddison et al., 2017; Jang et al., 2017) to learn a state representation with binary latent vectors.
Their method is superior to model-free and other model-based approaches on the VizDoom 3D navigation task.
Unlike our work, they do not focus on the multi-goal planning setting.
But, their method is able to adapt to changes in the dynamics of the environment quickly.
Recent work by van der Pol et al.
(2020) involves learning a forward model with a continuous state representation using a loss function based on the theory of
MDP homomorphisms (a generalization of bisimulation).
This work diﬀers from our work in that the discrete model is not learned directly–it is obtained using a heuristic that samples a large number of discrete states from encodings of observed states and then prunes them.
Other works consider action abstractions, included in the MDP homomorphisms framework, for the purpose of generalizing plans across problem instances (Bonet et al., 2019) and performing policy search in large action spaces (Biza and Platt, 2019).
In the context of optimal control, Watter et al.
(2015) used variational inference to learn a locally linear model of non-linear dynamics, and Matsubara et al.
(2014) encoded a continuous state space into discrete clusters with both a ﬂat and a factored Hidden
Markov Model in order to reduce the computational costs of solving Kullback-Leibler control problems (Kappen et al., 2013).
Appendix B. Additional Experiments B.1. Column World
The purpose of this experiment is to compare our method to a model-based approximate bisimulation baseline in a simple discrete environment.
Column World is a grid world with 30 rows and 3 columns (Lehnert and Littman, 2018).
The agent can move left, right, top and down, and it receives a reward 1 for any action executed in the right column; otherwise, it gets 0 reward.
Hence, the agent only needs to know if it is in the left, middle or right column, as illustrated in Figure 1.
First, we train a deep Q-network on this task and use it to generate a dataset of transitions.
As a baseline, we train a neural network model to predict r t and s t +1 given ( s t , a t ).
We then ﬁnd a coarse approximate bisimulation for this model using a greedy algorithm from (Dean et al., 1997) with the approximation constant (cid:15) set to 0.5.
We compare it with our method trained with an HMM prior on Q π ( s t , a t ) predicted by the deep Q-network.
We represent each state as a discrete symbol and use fully-connected neural networks for all of our models.
See Appendix B.2 for details.
Figure 3 shows the purity and the size of the abstractions found by our method and the baseline as a function of dataset size.
We need a ground-truth abstraction to calculate the abstraction purity–in this case, it is the three-state abstraction shown in Figure 1 right.
We assign each ground state to an abstract state (Figure 2) and ﬁnd the most common ground-truth label for each abstract state.
The abstraction purity is the weighted average of the fraction of members of an abstract states that share its label.
We include a snippet of code that computes purity in Appendix B.1.
Both methods can ﬁnd an abstraction with high purity.
However, approximate bisimulation does not reduce the state space (there are 90 ground states) until the model of the environment is nearly perfect, which requires more than 11000 training examples.
Our method always ﬁnds an abstraction with six states (the number of abstract states is a hyper-parameter), but our method ﬁnds a compact high-purity abstraction much faster than the baseline.
Notice that we parameterize our method with more abstract states than the size of the coarsest bisimulation.
In practice, this over-parameterization aids convergence.
B.2. Shapes World GMM and HMM comparison
We test the ability of our algorithm to ﬁnd accurate bisimulation partitions.
Table 2 shows the results for our method for both the GMM and the HMM prior
Both models reach a high abstraction purity (described in Section B.1) in all cases except for the three objects stacking task in a 4 × 4 grid world.
The smallest MDP for which a bisimulation exists containsabstract states; our algorithm has 1000 possible abstract states available.
Our experiment shows that the HMM prior can leverage the temporal information, which is missing from the GMM, to allocate abstract states better.
Setting
GMM
HMM 2 pucks,
[Table]
grid
[Table]
[Table] [Table] [Table]
[Table] [Table]
[Table]
[Table] [Table]
[Table] [Table]
[Figure]
We vary the dataset sizes used for learning the bisimulation from 1000 to 20000 samples
Abstraction size refers to the number of abstract states.
Appendix C. Theoretical Analysis
We analyze the idealized case where the following assumptions hold:
The transitions in the ground MDP are deterministic; 2.
The HMM prior has enough components to represent a bisimulation, no two components share a mean; 3.
The prior over hidden states p ( c t = k ) is ﬁxed; 4.
The encoder is deterministic and the prior observation model is an identity function over component means; 5.
The decoder p ( y | z t , a ) makes a prediction for each component using a table of stateaction values.
Assumptions on the MDP
Let M = (cid:104) S, A, T, R, γ (cid:105) be a ground MDP with ﬁnite S and A , an arbitrary R and a deterministic T .
We assume there exists a bisimulation mapping φ bisim : S → C with K abstract states ( | C | = K ).
φ bisim does not have to be the coarsest bisimulation (i.e. the one using the lowest possible number of abstract states).
Assumptions on the model
Let f : S → R D be a deterministic encoder.
For simplicity, we use a Hidden Markov
Model with a Kronecker delta observation model parameterized by K component means
C = { c 1 , c 2 , ..., c K } with probability density function
[Equation]
The encoder and the Hidden Markov Model induce an abstraction mapping φ : S → C
[Equation]
We assume ∀ c,c (cid:48)
∈ C µ c (cid:54) = µ c (cid:48) (i.e. no two components share their means) and ∀ s ∈ S ∃ c
∈ C
f ( s ) =
µ c (i.e. each encoding equals to some cluster mean).
The component prior p ( c ) assign a uniform probability to each cluster to match our experimental setting.
Finally, we assume the decoder p ( y | z, a ) is tied to the cluster assignment φ .
That is, each components has a parameter q c ∈ R | A | , which stores the state-action value of this component for each action.
The decoder probability density is a normal distribution with a mean of q c,a given an action a and a ﬁxed standard deviation σ
[Equation]
[Equation]
We can decompose the objective into a Q-value loss L Q and a prior loss L P
L Q
[Equation] [Equation]
Deﬁnition
An abstraction mapping φ is homogeneous if for each s, s (cid:48) ∈ S, a ∈ A, ˆ c ∈ C φ ( s )
=
φ ( s (cid:48) ) implies (cid:80) ˆ s ∈ φ − 1 (ˆ c ) T ( s, a, ˆ s )
= (cid:80) ˆ s ∈ φ − 1 (ˆ c ) T ( s (cid:48) , a, ˆ s ) .
Lemma 2 Let φ be an abstraction mapping that is a Q ∗ -irrelevance abstraction and is also homogeneous.
Then φ is a bisimulation.
Proof Fix an abstract state c and an action a .
For each s, s (cid:48) ∈ φ − 1 ( c ) we have that Q ∗ ( s, a ) =
Q ∗ ( s (cid:48) , a ) (by Q ∗ -irrelevance) and ∀ ˆ c ∈ C (cid:80) ˆ s ∈ φ − 1 (ˆ c ) T ( s, a, ˆ s ) =
(cid:80) ˆ s ∈ φ − 1 (ˆ c ) T ( s (cid:48) , a, ˆ s ) (by φ being homogeneous).
Let us expand the state-action values of ( s, a ) using the Bellman equation
[Equation]
We can perform a similar expansion for ( s (cid:48) , a )
[Equation]
We write T ( s, a, ˆ c ) as a shorthand for (cid:80) ˆ s ∈ φ − 1 (ˆ c ) T ( s, a, ˆ s ).
T ( s (cid:48) , a, ˆ c ) =
T ( s, a, ˆ c ) holds for all abstract states
if φ ( s ) =
φ ( s (cid:48) ) because φ is homogeneous.
All states mapped to a particular abstract state have the same value (for an optimal policy) due to Q ∗ -irrelevance; hence, we can write V ∗ (ˆ c ).
Since Q ∗ ( s, a ) =
Q ∗ ( s (cid:48) , a ) and x = x it follows that R ( s, a )
=
R ( s (cid:48) , a ) (i.e. φ is reward-respecting).
A reward-respecting homogeneous state abstraction is a bisimulation by deﬁnition.
Lemma 3 There exist model parameters θ such that L Q ( θ ) = 0 .
Proof Strategy: we can achieve zero L Q by assigning states into components such that states in each component have equal state-action values for all actions.
We need to show that there are enough components to perform this assignment.
Parameters θ induce an abstraction mapping φ .
Assume that L Q ( θ ) = 0.
Then for each c ∈ C , s, s (cid:48) ∈
φ − 1 ( c ) , a ∈ A we have ( Q ∗ ( s, a ) − Q ∗ ( s (cid:48) , a )) 2 = 0, which implies | Q ∗ ( s, a ) − Q ∗ ( s (cid:48) , a ) | = 0.
Hence, φ is a Q ∗ -irrelevance abstraction (Li et al., 2006).
By
Li et al.
(2006), for each bisimulation abstraction there exists a Q ∗ -irrelevance abstraction that is equal in size or coarser (i.e. it uses fewer abstract states).
By our assumption that we have enough components to represent a bisimulation abstraction, there are also enough components to represent a Q ∗ -irrelevance abstraction.
Lemma 4 There exist model parameters θ such that L P ( θ ) =
log K .
L P ( θ ) =
log K is the global minimum.
Proof Under our assumptions, we can reduce the prior loss function to
[Equation]
The minimum of this loss function is achieved when the term p ( φ ( f ( s t +1 ))
| φ ( f ( s t )) , a ) =for all states and actions.
Since we assume the transition dynamics of the ground MDP are deterministic, this transition model is only possible if the abstraction mapping φ is homogeneous.
Our model can represent such abstraction mapping because, by our assumption, it can represent a bisimulation, which is homogeneous.
Theorem 5 There exist model parameters θ that reach the global maximum of L IB ( θ ) =
− β log K with β > 0 .
The abstraction mapping φ induced by any such model parameters is a bisimulation.
Proof Since we assume we have enough components to represent a bisimulation, we can represent a Q ∗ -irrelevance abstraction (by Lemma 2) that is also homogeneous (by Lemma 3).
Any such abstraction is a bisimulation by Lemma 1.
Appendix D. Experimental Details D.1. Columns World
The deep Q-network that is used to collect the dataset has two hidden layers of 256 neurons followed by ReLU activation functions.
We train it for 40000 time steps with an (cid:15) -greedy policy; (cid:15) linearly decays from 1 to 0.1 over 20000 time steps.
We use a learning rate of 0.0001, 32 mini-batch size, the target network is updated every 100 time steps and we use prioritized replay with the default settings (Schaul et al., 2016).
The optimizer used for training is mini-batch gradient descent with momentum set to 0 .
The dataset for training the abstract and direct models is collected after training with (cid:15) set to 0.5.
We compute the abstraction purity over every possible ground state.
Each state is represented as a 90-dimensional one-hot encoded vector.
As a baseline, we train a model with two fully-connected layers of 128 neurons followed by ReLU activations and two heads, one for predicting the reward and the other for the next state.
We use a mean squared error loss for the reward prediction and cross-entropy loss for the next state (we treat each dimension of the predicted 90-dimensional vector as a probability of being in that particular state).
Finally, we run an approximate partition iteration algorithm following
Dean et al.
Our model with an HMM prior uses the same architecture as the above model, except it makes only one prediction: the state-action value associated with a given state-action pair.
We set the number of hidden states to 6, the observation model of the HMM is 32-dimensional, encoder and model learning rates are 0 .
01, β is 0.0001, the means of the HMM observations are initialized with 0 mean and 0.01 standard deviation and the diagonal covariances are initialized with -1 mean and 0.1 standard deviation before being exponentiated.
We train the models using the Adam optimizer (Kingma and Ba, 2015).
D.2. Shapes World
For dataset collection, the input image is resized to 64 × 64 before being fed into a deep
Q-network.
We use four convolutional layers with 32, 64, 128, and 256 ﬁlters; the ﬁlter size is four and the stride is set to two (each convolutional downsamples the input by a factor of two); we use ”same” padding.
The convolutions are followed by a single fully-connected layer with 512 neurons and a head for predicting the state-action values.
The learning rate is set to 0.00005, the batch size is 32, the buﬀer size is 100000 and we train for 100000 steps.
Actions are selected with an (cid:15) -greedy policy– (cid:15) is linearly decayed from 1.0 to 0.1 overtime steps.
We collect a dataset of 100000 transitions after training the model with (cid:15) set to 0.1.
80% of the dataset is used for training and 20% for computing the abstraction purity.
Our model uses the same neural network, except we insert batch normalization between each layer and its activation function (we use ReLU) (Ioﬀe and Szegedy, 2015).
The model predicts a 32-dimensional vector of means and a diagonal covariance, from which we sample the continuous encoding z .
The GMM or HMM uses 1000 components (hidden states), the initial means of the components are drawn from a Gaussian distribution with 0 mean and 0.1 variance.
The variances are drawn from a Gaussian with -1.0 mean and 0.1 before being exponentiated.
We train the model for 50000 steps, then we collect batch normalization statistics over the whole dataset, and we resume training only the prior with a ﬁxed encoder and unfrozen component weights p ( c t ) (previously held uniform ﬁxed) for anothersteps.
β is set to 10 − 6 , encoder learning rate to 10 − 3 and prior learning rate to 10 − 2 .
We train the model with Adam optimizer (Kingma and Ba, 2015).
To get a reward function over the abstract MDP induced by the HMM, we ﬁnd abstract states with 99% of ground states that are mapped to them being goal states for a given goal.
We plan state-action values for each abstract-state action pair using Value Iteration and run an agent with a softmax policy with τ set to 10 − 2 for 100 episodes.
We ran a hyper-parameter search for the learning rates and β on the task of stacking three pucks and then used the same parameters for all other experiments.
We tried the following learning rates: { 0 . 005 , 0 .
001 , 0 . 0005 , 0 . 0001 } and β : { 0 .