Year: 2018

Abstract:
Exploration in environments with sparse rewards is a key challenge for reinforcement learning. How do we design agents with generic inductive biases so that they can explore in a consistent manner instead of just using local exploration schemes like epsilon-greedy? We propose an unsupervised reinforcement learning agent which learns a discrete pixel grouping model that preserves spatial geometry of the sensors and implicitly of the environment as well. We use this representation to derive geometric intrinsic reward functions, like centroid coordinates and area, and learn policies to control each one of them with off-policy learning. These policies form a basis set of behaviors (options) which allows us explore in a consistent way and use them in a hierarchical reinforcement learning setup to solve for extrinsically defined rewards. We show that our approach can scale to a variety of domains with competitive performance, including navigation in 3D environments and Atari games with sparse rewards.

heading: 1 INTRODUCTION
text: Exploration in environments with sparse feedback is a key challenge for deep reinforcement learning (DRL) research. In DRL, agents typically explore with local exploration strategies like epsilon greedy or entropy based schemes. We are interested in learning structured exploration algorithms, grounded in spatio-temporal visual abstractions given raw pixels. In human perception and its developmental trajectory, spatio-temporal pixel groupings is one of the first visual abstractions to emerge, which is also used for intrinsically motivated goal-driven behaviors (Spelke & Kinzler, 2007). Inspired by this insight, we develop a new agent architecture and loss functions to autonomously learn visual abstractions and ground temporally extended behaviors in them.
Our approach and key contributions can be broken down into two parts: (1) an information theoretic loss function and a neural network architecture to learn visual groupings (abstractions) given raw pixels and actions, (2) a hierarchical action-value function agent which explores in the space of options grounded in the learned visual abstractions, instead of low level actions.
In the first step, we pass images through an encoder which outputs spatial discrete vector-quantized (VQ‘) grids, with 1 of E discrete components. We train this encoder to maximize the mutual information between VQ layers at different time steps, in order to obtain a temporally consistent representation, that preserves controllability and appearance information. We extract segmentation masks from the VQ layers for the second step, referred to as visual entities. We compute affine geometric measurements for each entity, namely centroid and area of the corresponding segment. We use off-policy learning to train action-value function to minimize or maximize these measurements, referred collectively as the options bank. Controlling these measurements enable higher levels of behaviors such as approaching an object (maximizing area), avoiding objects (minimize area or minimize/maximize centroid coordinates), moving an object away towards the left (minimize centroid x coordinate), controlling the avatars position on the screen etc.
Finally, given a task reward, we use off-policy learning to train a meta action-value function that takes actions at fixed intervals and selects either one of the policies in the options bank or low-level actions. So effectively, this hierarchical action-value function setup solves a semi markov decision process as in (Sutton et al., 1999; Kulkarni et al., 2016).
We demonstrate that our approach can scale to two different domains – navigation in a 3D environment and challenging Atari games – given raw pixels. Although much work remains in improving the visual and temporal abstraction discovery models, our results indicate that it is possible to learn bottom-up structured exploration schemes with simple spatial inductive biases and loss functions.

heading: 2 RELATED WORK
text: Learning visual abstractions has a long history in computer vision with some of the earlier successes relying on clustering either for inference (Shi & Malik, 2000) or learning (Ren & Malik, 2003). More recently some of these intuitions were adapted to neural networks as in (Xia & Kulis, 2017). Instance segmentation algorithms have been developed to output spatio-temporal groupings of pixels from raw videos (Romera-Paredes & Torr, 2016). However, most of the existing deep learning based approaches require supervised data. Structured deep generative models (van den Oord et al., 2017) is another approach to learning disentangled representations from raw videos. Very recently segmentation has been cast as a mutual information maximization problem in (Ji et al., 2018) where the mutual information is computed between the original image segmentation and the output of its transformation (additive color changes in HSV space, horizontal flips and spatial shifts). That approach uses the discrete mutual information estimate which makes it only applicable to enforce pixel-label constraints. For continuous variables recent papers have proposed very promising techniques. Most relevant to our work are (van den Oord et al., 2018) and (Belghazi et al., 2018).
In reinforcement learning research, semi-MDPs and the options framework (Sutton et al., 1999) have been proposed as a model for temporal abstractions of behaviors. Our work is most similar to hierarchical-DQN (Kulkarni et al., 2016). However, this approach required hand-crafted instance segmentation and the agent architecture is not distributed to learn about many intrinsic rewards learners at the same time. Object-Oriented-MDPs (Diuk et al., 2008) uses object oriented representations for structured exploration but requires prebuilt symbolic representations. A recent paper also demonstrates the importance of object based exploration when humans learn to play video games (Dubey et al., 2018). HRA (Van Seijen et al., 2017) is an agent that used prebuilt object representations to obtain state of the art policies on Pacman using object based structured exploration. Another interesting line of work is (Gregor et al., 2016) which formalizes a notion of empowerment which ends up as a mutual information between options in the same MDP. Count-based exploration algorithms have yielded impressive results on hard exploration Atari games (Ostrovski et al., 2017).
The Horde architecture (Sutton et al., 2011) proposed learning many value functions, termed as Generalized Value Functions (GVFs), using off-policy learning. This work was later extended with neural networks by Schaul et al. (2015). Our approach automatically constructs GVFs or a UVFA using the abstract entity based representations. Our work is also related to pixel control (Jaderberg et al., 2016) as an auxiliary task. However, we learn to control and compose abstract discrete representations.
Notation. Unless explicitly stated otherwise, we will use calligraphic upper-case letters for sets and MDPs, upper-case letters for random variables, bold capitals for constants and lower-case letters for realizations/measurements.

heading: 3 MODEL
text: Consider a Markov Decision ProcessM = (S,A,P, r) (MDP) represented by states s ∈ S, actions a ∈ A, transition distribution function P : S × A × S → [0, 1] and an extrinsic reward function defined as r : S × A → R. In a discrete MDP, an agent observes a state st at time t, produces and action at then the agent can observe st+1 ∼ P(S′|S = st, A = at) and a reward rt+1 = r(st, at). The agent’s objective is to maximize the expected sum of rewards over time. In this work we are focusing on visual inputs thus we assume S ⊂ RH×W, where H,W are the height and width of the image. This is a very important special case in current applications but many of the intuitions and machinery we develop should carry over to different domains.
Our agent also learns an separate abstract representation we call visual entities1 and denote with v ∈ V , where V ⊂ {1 . . .E}H×W i.e. it assigns an entity id to each location. These are meant to capture useful information about the state s and form the basis for computing the intrinsic rewards re,m (see Sec. 3.1). Here e ∈ {1 . . .E} denotes a discrete entity id, where E is the maximum number of possible entities, and m ∈ {1 . . .M} denotes a geometric feature of e’s derived segmentation mask, from M possible measurements. In this work the measurements are fixed to what we consider a sufficient set that captures the essential information for natural 3D navigation and Atari game play, and settled on centroid cartesian coordinates and entity mask area, which we can both minimize and maximize, thus in all our experiments M = 6. Temporal changes in these measurements constitute our intrinsic reward functions2 which induce E ×M additional MDPs Oe,m = (S,A,P, re,m). These intrisic rewards will induce behaviors which should hopefully provide structured exploration in the original MDP i.e. picking a random Qe,m is more likely to lead to higher reward than epsilon greedy exploration.
Our agent architecture tries to leverage this additional structure. The top level MDP M is represented by Qmeta which outputs action ametat at time t (switched every T time steps), where the action space is discrete 1 of (E ×M) + 1 possible actions. In our implementation this is modelled by composite actions E + 1 and M i.e. ametat = (et,mt) with et ∈ {1 . . .E + 1} and mt ∈ {1 . . .M}. We also learn (E ×M) + 1 separate Q functions: (E ×M) that each solve one of the Oe,m and one for the original MDP Qtask. These Q functions are defined over the environment action set and differ only in the reward function (see Figure 2 for schematic representation).

heading: 3.1 VISUAL ABSTRACTIONS
text: Our agent relies on an abstraction model that assigns each pixel in the image to one of E separate abstractions or entities. To obtain this representation, the image is passed through a convolutional network (CNN) encoder to output a spatial grid of the same resolution as the original image. Then a vector quantized layer V (van den Oord et al., 2017) assigns the planar activations to 1 of E entities. From the agents perspective this means all the pixels that are being grouped together become indistinguishable thus providing the visual abstraction we desire.
Let us define f : RH×W → {1 . . .E}H×W the function that takes the observation st at time t and computes an abstract representation corresponding to it vt = f(st). The key question is how to train f . One way is to make it representative of the current state. In most of the current literature, this is
1We use visual entities and visual abstractions interchangeably. Both terms capture the most important qualities of the representation: 1) emergent nature and 2) discreteness.
2If the mask is empty then the centroid is (0, 0) and area is zero.
achieved by training a function g : {1 . . .E}H×W → RH×W such that g(f(st)) ' st. This ensures that the representation preserves all the information about the state which is a sufficient condition for the purpose. It is not hard to see however that it potentially wastes model capacity on unimportant factors of variation. In this work we aim instead for a necessary condition for constraining the representation. To this end we train f to be an injective function i.e. such that a decoding function g can distinguish between states by looking at the representation ∀x 6= y =⇒ g(f(x)) 6= g(f(y)). Our approach is to formulate classification losses to promote distinguishability between different states at different levels of the representation. This was shown, in works like (Barber & Agakov, 2004; Belghazi et al., 2018; van den Oord et al., 2018), to be equivalent to maximizing a lower-bound mutual information between random variables of the representation3. By choosing the appropriate random variables and sampling strategies from the joint distribution and marginals we can specify the right representation invariances as follows (see Figure 1 for an overview):
Preserving global information. The main term driving the representation learning is a global information term. To estimate it, the VQ layer output at time t, vt, is further processed by another CNN encoder to output a frame level embedding vector Gt. We train a non-parametric classifier to distinguish pairs of frames from the same trajectory from pairs of frames from different trajectories. For that we form pairs (gt, gt+∆), where ∆ is sampled randomly, from pairs (gt, g′), where g′ is sampled randomly from a different trajectory. Training this classifier lets us indirectly maximize a lower-bound on I(Gt, Gt+∆). This forces V to preserve enough information that distinguishes this particular trajectory from other ones which tends to remove all irrelevant information like textures and unchanging “background” elements and preserve useful moving elements. Note that as long as predictions are stable across time there is no pressure exerted by this cost to simplify the representation.
3We write our losses as I because of this fact but this is an approximation.
Preserving controllable information. Secondly, we want the controllable information to be preserved in the abstract representation. That is, we want to know what aspects of the input were changed as a result of the action. We achieve this by training to predict which action was taken in a particular transition based on our representation. For that we add another loss, denoted by I((Gt, Gt+1), At), that maximizes a lower-bound on mutual information between the a pair of consecutive frames and the action that was taken in the transition.
Preserving local appearance information. Finally, for hard exploration Atari games, where inputs change very little e.g. Montezuma’s Revenge much of the initial experience is from the first room, we add an appearance term4. This term is meant to align the abstract representation with appearance changes i.e. promote abstractions have consistent colors. For that we feed the input image into a shallow CNN encoder which outputs an embedding C of the local color and texture structure of the image. We would like for the abstract representation to follow the appearance changes thus we maximize a lower bound on mutual information I(Vt, Ct) in the same way as before. The positive pairs spatially aligned V and C embeddings, and the negative pairs are obtained by sampling appearance embeddings from other spatial locations and pairing them with non-sampled V . This promotes V representations that represent well appearance under the representational constraints of the VQ representation.
To learn the agents’ abstract representation we minimize a weighted sum of these classification losses
Labs(θabs) = ∑ t −αgEg0∼P(G′|gt) gi∼P(G) log qg(Zg = 0|(g0, . . . , gK), gt) (1)
−αc log qc(at|(gt, gt+1))− αaEcit∼P(Ct) vt∼P(Vt) log qa(Zc = 0|cit, vt) (2)
where we denote by lower case letter the samples from the corresponding uppercase random variables. P(G) is the time independent marginal distribution over Gt’s. We model qg as
qg(Zg = 0|(g0, . . . gK), gt) = expφ(g0, gt)∑K i=0 expφ(gi, gt)
(3)
with φ(., .) being the cosine similarity over the embeddings andK is a hyper parameter denoting the number of samples. We can similarly derive the third term. This is possible because the embeddings have the same dimensionality. The second term is a cross entropy based action classifier.

heading: 3.2 BEHAVIOR ABSTRACTION
text: The agent is represented primarily by three sets of Q functions: Qmeta, Qtask and {Q1,1, . . . Qe,m} from the options bank. Each Q function has a corresponding policy we denote by π, in our experiments these are epsilon greedy policies with respect to the corresponding Q function. We denote T to be the fixed temporal commitment window for Qmeta, which means that it acts every T steps. Note that Qtask and Qe,m act at each environment time step. We can express all three Q functions as: (1) Qmeta(s, (e,m)) = E [∑∞ t′=t γ t′−trt′ |st = s, ametat = (e,m), πmeta ] , (2) Qtask(s, a) =
E [∑∞
t′=t γ t′−trt′ |st = s, ataskt = a, πtask
] and (3) Qe,m(s, a) = E [∑∞ t′=t γ
t′−tre,mt′ |st = s, at = a, πe,m ] .
We represent each Q function with a deep Q network (Mnih et al., 2015), for instance Qtask(s, a) ≈ Qtask(s, a; θtask). Each Q ∈ {Qtask, Qmeta, Q1,1, ..., Qe,m} can be trained by minimizing corresponding loss functions – Ltask(θtask),Lmeta(θmeta) and {L1,1bank(θ1,1), ...,L e,m bank(θ
e,m)}. We store experiences (st, (et,mt, at), rt, st+1) in D, a shared buffer from which all the different Q functions can sample to perform their updates. The transitions are stored in such a way as to be able to sample trajectories.
The Q-learning (Sutton & Barto, 1998) objective can be written as minimizing the loss:
Ltask(θtask) = E(sτ ,aτ ,rτ ,sτ+1)τ∼D[(Rτ −Q task(sτ , aτ ; θtask)) 2] (4)
4We found that for tasks that have richer visual variation this was not necessary.
where Rτ = ∑τ+U−1 t=τ γ t+τrt + γ t+U maxa′ Q task(sτ+U, a ′; θtask), where U is the length of the unroll. The loss function Lmeta and all the Le,mbank can be written in a similar fashion. In our experiments we use Q(λ) and learn all parameters with stochastic gradient descent by sampling experience trajectories from the shared replay buffer (see Algorithm 1 for details).

heading: 4 EXPERIMENTS
text: The implementation has a learning setup inspired by the batched actor critic agent (Espeholt et al., 2018), with a central GPU learner and multiple actors (64 in most experiments). See 6 for details regarding visual abstraction architecture and agent network architecture.
The exploration setup. For the baseline we have 3 types of actors differing only in the exploration parameter corresponding to high exploitation, high exploration and medium exploration. From the total number of 64 actors that corresponds to 20, 10 and 34 respectively with epsilon values of .001, .5, .01. Both the meta agent and base policies require exploration so we keep the same split but the high exploration is considered independently for the base and meta policies i.e. half the actors have meta = .5 and base = .01 and half have meta = .1 and base = .5. Furthermore we split equally the medium exploration actors in 3 groups with ( meta = .01, base = .1), ( meta = .33, base = .001) and ( meta = .01, base = .1) respectively. This provides much more stable learning at both the higher and lower levels of the behavior hierarchy.

heading: 4.1 NAVIGATION
text: We tested our approach on a 3D navigation domain with sparse rewards and hard exploration requirements. The domain, whose top down view (not visible to the agent) is shown in figure 4, contains four rooms each having a textured number at the entrance. The agent receives the image observation as well as hint, a number from 1 to 4, which indicates the number of the room which contains the target object, a green sphere. The goal is to reach the target as often as possible in the
Algorithm 1 Learning algorithm 1: Inputs: N number of episodes, base and meta exploration parameters, T the commitment
length, λtask, λmeta, λbank cost function parameters 2: Initialize experience replay buffer D and parameters {θmeta, θtask, θ1,1, ..., θe,m} for the meta-
control agent, task agent and options models respectively. 3: for i = 1 to N do 4: Initialize environment and get start state s 5: s0 ← s 6: while st is not terminal do 7: if t ≡ 0 mod T then 8: ametat = (et,mt)← EPSGREEDY(st, meta, θmeta) 9: end if
10: Compute abstract features vt from st (Section 3.1). 11: Compute intrinsic rewards rint = (re,m| ∀e ≤ E,m ≤M) from vt and vt−1 12: if et ≤ E then 13: at ← EPSGREEDY(st, base, Qet,mt) % selected Q function from bank 14: else 15: at ← EPSGREEDY(st, base, Qtask) % selected Qtask 16: end if 17: Execute at and obtain next state st+1 and extrinsic reward rt from environment 18: Store transition (st, (et,mt, at), rt, st+1) in D 19: Use RMSProp to optimize Labs(θabs) %see eq. (2)) 20: Use RMSProp to optimize λtaskLtask + λmetaLmeta + λbank 1ME ∑E e=1 ∑M m=1 L e,m bank
(see eq. (4)) 21: end while 22: end for
limited time budget (250 steps). Once the goal is reached the target is relocated in another randomly chosen room and the hint is updated accordingly. To solve the task the agent needs to associate the hint with its position in the environment, which in turn means exploring far away regions in the maze because the target may move in a far away place. The policy bank provides exactly this type of exploration basis set. A baseline agent with the same architecture but without the policy bank cannot solve this task in 100M steps see Figure 4. In this domain we can also query the environment representation to determine if a given pixel comes from either a wall, the skyline, floor, one of the textured numbers or the target sphere. An agent with our proposed architecture but computing intrinsic rewards based on this priviledged information solves the tasks easily. Most interestingly, an agent that uses the abstraction inference method we propose can also solve it too, though with slightly worse performance.
To gain an insight into the evolution of the agent policy, we plot in Figure 4 (c), a histogram of the Qmeta policy actions during training (time is on the vertical axes and flows from top to bottom). The leftmost 20 actions represent the policy bank and the rightmost one represents Qtask. We can see that the agent uses most of the policies in the bank but as training progresses learns to rely more and more on Qtask which is the optimal behavior.

heading: 4.2 ATARI
text: We have run our agent with the same architecture and parameters on hard exploration Atari games where we can show that our agent has better performance than the baseline with a comparable architecture and losses (see Figure 5).

heading: 4.3 DM LAB
text: In order to see if our model scales to more visually challenging environments we have run our agent on 3 varied “DMLab-30” levels(Beattie et al., 2016). We show the training curves as a function of time in Figure 6. In the “non-match” task the agent teleports first to a room with one object then a room with two, one of which is the same as before. To get a positive reward the agent has to move on top of the other object, otherwise it receives a negative reward. The sequence continues for a fixed number of steps. Humans achieve 66 points and a state of the art agent gets 26. Our Q learner baseline achieves only 9 points whereas our proposed agent achieves 33. Though this is meant to be a memory task structured exploration seems to help achieve much better scores than the baselines.
We have also considered an experiment on the “keys doors” task. In this task the agent has to successively pickup keys to unlock doors which seems like a task structure were our representation could be more effective than the baseline. We found that, though both methods are competitive, our representation was not sufficent to learn a better policy. We think that this may be due to noise in the abstraction inference as well as a planning aspect that is not well enough handled in our agent in its current form. Finally, our agent is outperformed by the baseline on the challenging watermaze task. This is most likely due to the policy bank exploring mostly straight trajectories rather than circular ones which are more appropriate on this task.

heading: 5 DISCUSSION
text: We have shown that it is possible to design unsupervised structured exploration schemes for modelfree DRL agents, with competitive performance on a range of environments given just raw pixels.
One of the biggest open question moving forward is to find strategies to balance structure or inductive biases and performance. Our current solution was to augment the meta-controller with Qtask along with the options bank as sub-behaviors. The typical strategy that agents follow is to rely on the options bank early in training and then use this experience to train the Qtask policy for optimality as training progresses. This is reasonable given that the options models may not cover the optimal policy but could serve as a good exploration algorithm throughout training. As new unsupervised architectures and losses are discovered, we expect to narrow the gap between the optimal desired behaviors and the options bank.
Learning visual entities from pixels is still a challenging open problem in unsupervised learning and computer vision. We expect novel sampling schemes in our proposed architecture to improve the entity discovery results. Other unsupervised video segmentation algorithms and discrete latent variable models could also be used to boost the discovery process.

heading: 6 SUPPLEMENTARY MATERIAL
text: Visual Abstraction Architecture. The encoder for the abstractions is a set of 3 convolutional layers 3×3 kernels with 64 features each followed by ReLU nonlinearities. After a 1×1 convolutional layer with 8 outputs the features are l2 normalized. We found that to work better in practice with the VQ layer sitting on top. The VQ layer has 8 elements unless otherwise stated. Since the strides are always 1 the VQ output has the same spatial dimensions as the input image. The global loss is a stack of convolutions 3 × 3, followed by a 2 × 2 max pooling with stride two to reduce the resolution of the output and then a ReLU non-linearity. The output is then flattened to give the global embedding of the abstract representation. The image embedding is a two-layer 8 filter convnet with a ReLU non-linearity inbetween.
Agent architecture. The agent architecture is a standard 3 layer convolutional stack similar to Mnih et al. (2015) with 512 hidden unit output and an LSTM on top. The output of the LSTM is fed into an visual abstraction selection Q function, a measurement selection Q function, a regular task Q function layer and the policy bank i.e. a layer with M × E × num actions outputs. The Q function layers are all dueling heads as in Wang et al. (2016).
Setup and baseline. Our setup is very similar to the one in (Espeholt et al., 2018). Multiple actors (64 in most examples) in the acting loop that send trajectories to a shared learner which processes them in batched fashion (batch size of 32). The main difference is the use of value based Q(λ) loss (with a λ value of .85) instead of actor-critic with off-policy correction. The baseline agent has the same exact architecture and loss as our agent. In fact if we ignored the meta control Q function and the options bank we get our baseline exactly.

Review summaries:
The appproach introduces visual abstractions that are used for reinforcement learning. The abstractions are learned using a lower bound on the mutual information and options are created to generate different measurements for each abstraction. The algorithm hence learns to "control" each abstraction as well as to select the options to achieve the overall task. The algorithm is tested on a 3D navigation task and a few Atari tasks which are known for difficult exploration.