Deep Successor Reinforcement Learning

Introduction

Many learning problems involve inferring properties of temporally extended sequences given an objective function. For instance, in reinforcement learning (RL), the task is to find a policy that maximizes expected future discounted rewards (value). RL algorithms fall into two main classes: (1) model-free algorithms that learn cached value functions directly from sample trajectories, and (2) model-based algorithms that estimate transition and reward functions, from which values can be computed using tree-search or dynamic programming. However, there is a third class, based on the successor representation (SR), that factors the value function into a predictive representation and a reward function. Specifically, the value function at a state can be expressed as the dot product between the vector of expected discounted future state occupancies and the immediate reward in each of those successor states.

Representing the value function using the SR has several appealing properties. It combines computational efficiency comparable to model-free algorithms with some of the flexibility of model-based algorithms. In particular, the SR can adapt quickly to changes in distal reward, unlike model-free algorithms. In this paper, we also highlight a feature of the SR that has been less well-investigated: the ability to extract bottleneck states (candidate subgoals) from the successor representation under a random policy . These subgoals can then be used within a hierarchical RL framework. In this paper we develop a powerful function approximation algorithm and architecture for the SR using a deep neural network, which we call Deep Successor Reinforcement Learning (DSR). This enables learning the SR and reward function from raw sensory observations with end-to-end training.

The DSR consists of two sub-components: (1) a reward feature learning component, constructed as a deep neural network, predicts intrinsic and extrinsic rewards to learn useful features from raw observations; and (2) an SR component, constructed as a separate deep neural network, that estimates the expected future “feature occupancy” conditioned on the current state and averaged over all actions. The value function can then be estimated as the dot product between these two factored representations. We train DSR by sampling experience trajectories (state, next-state, action and reward) from an experience replay memory and apply stochastic gradient descent to optimize model parameters. To avoid instability in the learning algorithm, we interleave training of the successor and reward components.

We show the efficacy of our approach on two different domains: (1) learning to solve goals in grid-world domains using the MazeBase game engine and (2) learning to navigate a 3D maze to gather a resource using the Doom game engine. We show the empirical convergence results on several policy learning problems as well as sensitivity of the value estimator given distal reward changes. We also demonstrate the possibility of extracting plausible subgoals for hierarchical RL by performing normalized-cuts on the SR .

Related work

The SR has been used in neuroscience as a model for describing different cognitive phenomena. showed that the temporal context model , a model of episodic memory, is in fact estimating the SR using the temporal difference algorithm. introduced a model based on SR for preplay and rapid path planning in the CA3 region of the hippocampus. They interpret the SR as an an attractor network in a low–dimensional space and show that if the network is stimulated with a goal location it can generate a path to the goal. suggested a model for tying the problems of navigation and reward maximization in the brain. They claimed that the brain’s spatial representations are designed to support the reward maximization problem (RL); they showed the behavior of the place cells and grid cells can be explained by finding the optimal spatial representation that can support RL. Based on their model they proposed a way for identifying reasonable subgoals from the spectral features of the SR. Other work (see for instance, ) have also discussed utilizing the SR for subgoal and option discovery.

There are also models similar to the SR that have been been applied to other RL-related domains. introduced a model for evaluating the positions in the game of Go; the model is reminiscent of SR as it predicts the fate of every position of the board instead of the overall game score. Another reward-independent model, universal option model (UOM), proposed in , uses state occupancy function to build a general model of options. They proved that UOM of an option, given a reward function, can construct a traditional option model. There has also been a lot of work on option discovery in the tabular setting . In more recent work, Machado et al. presented an option discovery algorithm where the agent is encouraged to explore regions that were previously out of reach. However, option discovery where non-linear state approximations are required is still an open problem.

Our model is also related to the literature on value function approximation using deep neural networks. The deep-Q learning model and its variants (e.g., ) have been successful in learning Q-value functions from high-dimensional complex input states.

Model

where, $s_{t}$ is the state visited at time $t$ and the expectation is with respect to the policy and transition distribution. The agent’s goal is to find the optimal policy $Q^{*}$ which follows the Bellman equation:

2 The successor representation

The SR can be used for calculating the Q-value function as follows. Given a state $s$, action $a$ and future states $s^{\prime}$, SR is defined as the expected discounted future state occupancy:

Given the SR, the Q-value for selecting action $a$ in state $s$ can be expressed as the inner product of the immediate reward and the SR :

3 Deep successor representation

The SR for the optimal policy in the non-linear function approximation case can then be obtained from the following Bellman equation:

where $a^{\prime}=\text{argmax}_{a}m_{s_{t+1}a}\cdot\textbf{w}$.

4 Learning

The loss function for $\alpha$ is given by:

where $a^{\prime}=\text{argmax}_{a}u_{\alpha}(\phi_{s_{t+1}},a)\cdot\textbf{w}$ and the parameter $\alpha_{prev}$ denotes a previously cached parameter value, set periodically to $\alpha$. This is essential for stable Q-learning with function approximations (see ).

For learning $\mathbf{w}$, the weights for the reward approximation function, we use the following squared loss function:

Parameter $\theta$ is used for obtaining the $\phi(s)$, the shared feature representation for both reward prediction and SR approximation. An ideal $\phi(s)$ should be: 1) a good predictor for the immediate reward for that state and 2) a good discriminator for the states. The first condition can be handled by minimizing loss function $L^{r}_{t}$; however, we also need a loss function to help in the second condition. To this end, we use a deep convolutional auto-encoder to reconstruct images under an L2 loss function. This dense feedback signal can be interpreted as an intrinsic reward function. The loss function can be stated as:

The composite loss function is the sum of the three loss functions given above:

Automatic Subgoal Extraction

Learning policies given sparse or delayed rewards is a significant challenge for current reinforcement learning algorithms. This is mainly due to inefficient exploration schemes such as $\epsilon-$greedy. Existing methods like Boltzmann exploration and Thomson sampling offer significant improvements over $\epsilon$-greedy, but are limited due to the underlying models functioning at the level of basic actions. Hierarchical reinforcement learning algorithms such as the options framework provide a flexible framework to create temporal abstractions, which will enable exploration at different time-scales. The agent will learn options to reach the subgoals which can be used for intrinsic motivation. In the context of hierarchical RL, discuss a framework for subgoal extraction using the structural aspects of a learned policy model. Inspired by previous work in subgoal discovery from state trajectories and the tabular SR , we use the learned SR to generate plausible subgoal candidates.

Given a random policy $\pi_{r}$ ($\epsilon=1$), we train the DSR until convergence and collect the SR for a large number of states $\mathcal{T}=\{m_{s_{1},a_{1}},m_{s_{2},a_{2}},...,m_{s_{n},a_{n}}\}$. Following , we generate an affinity matrix $W$ given $\mathcal{T}$, by applying a radial basis function (with Euclidean distance metric) for each pairwise entry $(m_{s_{i},a_{i}},m_{s_{j},a_{j}})$ in $\mathcal{T}$ (to generate $w_{ij}$). Let $D$ be a diagonal matrix with $D(i,i)=\sum_{j}w_{ij}$. Then as per , the second largest eigenvalue of the matrix $D^{-1}(D-W)$ gives an approximation of the minimum normalized cut value of the partition of $\mathcal{T}$. The states that lie on the end-points of the cut are plausible subgoal candidates, as they provide a path between a community of state groups. Given randomly sampled $\mathcal{T}$ from $\pi_{r}$, we can collect statistics of how many times a particular state lies along the cut. We pick the top-k states as the subgoals. Our experiments indicate that it is possible to extract useful subgoals from the DSR.

Experiments

In this section, we demonstrate the properties of our approach on MazeBase , a grid-world environment, and the Doom game engine . In both environments, observations are presented as raw pixels to the agent. In the first experiment we show that our approach is comparable to DQN in two goal-reaching tasks. Next, we investigate the effect of modifying the distal reward on the initial Q-value. Finally, using normalized-cuts, we identify subgoals given the successor representations in the two environments.

We learn the optimal policy in the maze shown in Figure 2 using the DSR and compare its performance to the DQN . The cost of living or moving over water blocks is -0.5 and the reward value is 1. For this experiment, we set the discount rate to 0.99 and the learning rate to $2.5\cdot 10^{-4}$. We anneal the $\epsilon$ from 1 to 0.1 over 20k steps; furthermore, for training the reward branch, we anneal the number of samples that we use, from 4000 to 1 by a factor of 0.5 after each training episode. For all experiments, we prioritize the reward training by keeping a database of non-zero rewards and sampling randomly from the replay buffer with a 0.8 probability and 0.2 from the database. Figure 3 shows the average trajectory (over 5 runs) of the rewards obtained over 100k episodes. As the plot suggests, DSR performs on par with DQN.

Finding a goal in a 3D environment

We created a map with 4 rooms using the ViZDoom platform . The map is shown in Figure 2. We share the same network architecture as in the case of MazeBase. The agent is spawned inside a room, and can explore any of the other three rooms. The agent gets a per-step penalty of -0.01 and a positive reward of 1.0 after collecting an item from one of the room (highlighted in red in Figure2). As shown in Figure3, the agent is able to successfully navigate the environment to obtain the reward, and is competitive with DQN.

2 Value function sensitivity to distal reward changes

The decomposition of value function into SR and immediate reward prediction allows DSR to rapidly adapt to changes in the reward function. In order to probe this, we performed experiments to measure the adaptability of the value function to distal reward changes. Given the grid-world map in Figure2, we can train the agent to solve the goal specified in the map as highlighted in section 5.1. Without changing the goal location, we can change the reward scalar value upon reaching the goal from 1.0 to 3.0. Our hypothesis is that due to the SR-based value decomposition, our value estimate will converge to this change by just updating the reward weights $\mathbf{w}$ (SR remains same). As shown in Figure 4, we confirm that the DSR is able to quickly adapt to the new value function by just updating $\mathbf{w}$.

3 Extracting subgoals from the DSR

As shown in Figures 5 and 6, our subgoal extraction scheme is able to capture useful subgoals and clusters the environment into reasonable segments. Such a scheme can be ran periodically within a hierarchical reinforcement learning framework to aid exploration. One inherent limitation of this approach is that due to the random policy, the subgoal candidates are often quite noisy. Future work should address this limitation and provide statistically robust ways to extract plausible candidates. Additionally, the subgoal extraction algorithm should be non-parametric to handle flexible number of subgoals.

Conclusion

We presented the DSR, a novel deep reinforcement learning framework to learn goal-directed behavior given raw sensory observations. The DSR estimates the value function by taking the inner product between the SR and immediate reward predictions. This factorization of the value function gives rise to several appealing properties over existing deep reinforcement learning methods—namely increased sensitivity of the value function to distal reward changes and the possibility of extracting subgoals from the SR under a random policy.

For future work, we plan to combine the DSR with hierarchical reinforcement learning. Learning goal-directed behavior with sparse rewards is a fundamental challenge for existing reinforcement learning algorithms. The DSR can enable efficient exploration by periodically extracting subgoals, learning policies to satisfy these intrinsic goals (skills), and subsequently learning hierarchical policy over these subgoals in an options framework . One of the major issues with the DSR is learning discriminative features. In order to scale up our approach to more expressive environments, it will be crucial to combine various deep generative and self-supervised models with our approach. In addition to subgoals, using DSR for extracting other intrinsic motivation measures such as improvements to the predictive world model or mutual information is worth pursuing.

References