Title: QUANTILE REGRESSION REINFORCEMENT LEARNING WITH STATE ALIGNED VECTOR REWARDS

Abstract: Learning from a scalar reward in continuous action space environments is difficult and often requires millions if not billions of interactions. We introduce state aligned vector rewards, which are easily defined in metric state spaces and allow our deep reinforcement learning agent to tackle the curse of dimensionality. Our agent learns to map from action distributions to state change distributions implicitly defined in a quantile function neural network. We further introduce a new reinforcement learning technique inspired by quantile regression which does not limit agents to explicitly parameterized action distributions. Our results in high dimensional state spaces show that training with vector rewards allows our agent to learn multiple times faster than an agent training with scalar rewards.

1. INTRODUCTION:

Reinforcement learning (Sutton et al., 1998 ) is a powerful paradigm in which an agent learns about an environment through interaction. The common formulation consists of a Markov Decision Process (MDP) modeled as a 5-tuple (S, A, P, r, γ) where S is the (possibly infinite) set of states, A is the (possibly infinite) set of actions available to the agent, P : (S × A × S) → [0, 1] : P (s |s, a) is the transition probability of reaching state s ∈ S given state s ∈ S and action a ∈ A, r : (S × A) → R : r(s, a) is the reward received for taking action a in state s and γ is the reward discount factor. The goal of the agent is to maximize the cumulative discounted reward R = ∞ t=0 γ t r(s t , a t ) by choosing actions a t according to some (possibly stochastic) policy π : (S × A) → [0, 1] : π(a t |s t ). Sometimes it is further useful to make a distinction between the actual state space S and the correlated observation space O of the agent. In this case π : (O × A) → [0, 1] : π(a t |o t ) with o t ∈ O. The use of deep neural networks allowed this formulation to scale to high dimensional visual inputs approaching continuity in state space (Mnih et al., 2015) while others extended deep reinforcement learning to continuous action spaces (Lillicrap et al., 2015; Mnih et al., 2016) . While neural networks are powerful function approximators, they require large amounts of training data to converge. In the case of reinforcement learning this means interactions with the environment, a requirement easy to fulfill in simulation, yet impractical when the agent should interact with the real world. This problem is aggravated by the weak training signal of classical reinforcement learning -a simple scalar reward.

While originally the dopamine activity in mammal brains was linked to general "rewarding" events, Pinto & Lammel (2017) point out that the diversity of dopamine circuits in the mid brain is better modeled by viral vector strategies. Gershman et al. (2009) also show that human reinforcement learning incorporates effector specific value estimations to cope with the high dimensional action space. Inspired by these biological insights, we improve the sample efficiency of a deep reinforcement learning algorithm in this work by modeling a d-dimensional vector reward. A vector reward can in some domains easily be defined in alignment with the state space. We say that two vector spaces are aligned if their dimensions correlate and show that if state and action space are not aligned, a mapping from action distribution to state change distribution can be learned.

As a motivating example, consider the agent in Figure 1 (a) trying to reach the goal (marked by the blue dot). If we take p as the position vector of the agent relative to the goal, a sensible reward to guide the agent to the goal in this environment would be r = ||p|| -||p + a|| where || • || can be any norm in the vector space of the environment. For illustration purposes we'll focus on the L 1 norm in this example and throughout this paper. During training the agent might try action a which moves With a scalar reward this would be summarized as r = r x -r y , whereas a vector reward would keep the two reward dimensions distinguishable. (b) In most cases action and state space are however not aligned, therefore a mapping from action to state change must be learned. it closer to the goal in x direction, but a bit further away from the goal in y direction. The scalar reward would then just convey the information, that the action was rather positive (since the agent got closer to the goal) but miss out on the distinction that the action was good in x-direction but bad in y-direction. To provide this distinction, a more informative reward would keep the dimensions separate and therefore be a vector itself: r = |p| -|p + a| where | • | denotes the element-wise absolute value here. Note that this reward is dimension wise aligned with the position p, the state, of the agent. Since we focus on reaching problems in this work, we'll use the terms "position" and "state" interchangeably.

The problem with such a state aligned vector reward is however that the action space is in most cases not state aligned. To see this, consider the schematic robot arm in Figure 1 (b): The action dimensions a 1 and a 2 correspond to the torques of the robot arm and do not directly translate to a shift in x and y dimension, respectively. To address this issue we use the method proposed by Dabney et al. (2018b; a) to train a deep neural network to approximate the quantile function, in our case of the position change, given the current observation and quantile target. Additionally we give a parameterization of the action probability distribution as input to this position change prediction network (short PCPN). We then train the agent, parameterized by another neural network which maps from observations to action probability distributions, through a new reinforcement learning method we call quantile regression reinforcement learning (short QRRL). A schematic overview of our setup can be seen in Figure 2 .

To summarize, the contributions of this paper are the following:

• We extend the reinforcement learning paradigm to allow for faster training based on more informative state aligned vector rewards.

• We present an architecture that learns a probability distribution over possible state changes based on a probability distribution over possible actions.

• We introduce a new reinforcement learning algorithm to train stochastic continuous action policies with arbitrary action probability distributions.

2 QUANTILE REGRESSION AND IMPLICIT QUANTILE NETWORKS Quantile regression (Koenker, 2005) discusses approximation techniques for the inverse cumulative distribution function F -1 Y , i.e., the quantile function, of some probability distribution Y . Recent work (Dabney et al., 2018a; Ostrovski et al., 2018) shows that a neural network can learn to approximate the quantile function by mapping a uniformly sampled quantile target τ ∼ U([0, 1]) to its corresponding quantile function value F -1 Y (τ ) ∈ R. Thereby the trained neural network implicitly models the full probability distribution Y . 

W p (U, Y ) = 1 0 |F -1 Y (ω) -F -1 U (ω)| p dω 1/p

of distributions U and Y , also characterized as the L p metric of quantile functions (Müller, 1997) . Dabney et al. (2018b) show that the quantile regression loss (Koenker & Hallock, 2001) ρ

τ (δ) = (τ -1 δ<0 ) • δ (1)

minimizes the 1-Wasserstein distance of a scalar probability distribution Y to a uniform mixture of Diracs U . Here, δ = y -u with y ∼ Y and u ∼ U is the quantile sample error. Ostrovski et al. (2018) generalized this result by showing that the expected quantile loss

E τ ∼U ([0,1]) E z∼Z ρ τ (z -Qθ (τ )

) of a parameterized quantile function Qθ aproximating the quantile function F -1 Z of some distribution Z is equal to the quantile divergence

q(Z, Q θ ) := 1 0 F -1 Q θ (τ ) F -1 Z (τ ) (F Z (x) -τ )dx dτ

plus some constant not depending on the parameters θ. Here, Q θ is the distribution implicitly defined by Qθ . Therefore, training a neural network Qθ (τ ) to minimize ρ τ (z -Qθ (τ )) with z sampled from the target probability distribution Z effectively minimizes the quantile divergence q(Z, Q θ ) and thereby models an approximate distribution Q θ of Z implicitly in the network parameters θ of the neural network Qθ (τ ).

By approximating the quantile function instead of a parameterized probability distribution, as common in many deep learning models (Kingma & Welling, 2013; Lillicrap et al., 2015; Ha & Schmidhuber, 2018) , we do not enforce any constraint on the probability distribution Z, e.g., Z can be multi-modal, not continuous and non-Gaussian. This is crucial for our case as a position change distribution given a certain action distribution can have any shape.

In our setup, we train the PCPN with the quantile regression loss (1) to approximate a position change quantile function per position dimension. Aside from a target quantile τ ∼ U([0, 1]) per position dimension, the network input consists of an observation o ∈ O and a multi-variate Gaussian action probability distribution A = N (µ, Σ) with diagonal covariance matrix Σ = Iσ, parameterized by the mean µ and variance σ vectors. Therefore, the PCPN has the ability to implicitly learn the conditional position change distribution given an observation and an action distribution.

3. QUANTILE REGRESSION REINFORCEMENT LEARNING:

The dominant reinforcement learning algorithms are either value based methods, e.g., Q-learning (Watkins, 1989; Mnih et al., 2015) , policy gradient based methods, e.g., REINFORCE (Williams, 1992; Sutton et al., 1999) , or a combination of both, e.g., actor-critic methods (Sutton et al., 1998; Mnih et al., 2016) . In this section we establish a new reinforcement learning objective based on quantile regression. In contrast to Q-learning we allow for a continuous action space similar to policy based algorithms, but in contrast to policy based algorithms we do not limit ourselves to explicitly parameterized policies. Note that the technique described in this section is generally applicable to reinforcement learning problems and we therefore use the terms "action" and "policy" in the common reinforcement learning meaning, which is distinct from the meaning of "action" and "policy" in the rest of the paper. We later link in Section 4 the position change estimation of the PCPN to this new meaning of "action" to connect with the rest of the paper.

The main idea behind Quantile Regression Reinforcement Learning (or QRRL in short) is to model the policy implicitly in the network parameters, therefore allowing for complex stochastic policies which can find multi-modal stochastic optima in policy space. For this we model for each action dimension the quantile function

Qθ (τ, o) = F -1 Q θ (τ, o)

of the implicitly defined action distribution Q θ (o) by a neural network with an observation o and a target quantile τ ∈ [0, 1] as input. From the full network Qξ : [0, 1] d × O → R d , with d being the number of action dimensions, an action a ∈ A ≡ R d for a given observation o can be sampled by sampling τ ∼ U([0, 1] d ) and taking the network output as action. Since the network approximates quantile functions, the network output of a uniformly at random sampled quantile target is a sample from the implicitly defined action distribution. The question left to address is how to train the network, such that it (a) approximates the quantile functions of the action dimensions and (b) the implicitly defined policy maximizes the expected (discounted) reward R.

Here quantile regression comes in handy. Informally put, quantile regression is linked to the Wasserstein metric which is also sometimes refered to as earth movers distance. Imagine a pile of earth representing probability mass. In reinforcement learning we essentially want to move probability mass towards actions that were good and away from actions that were bad, where "good" and "bad" are measured by discounted accumulative (bootstrapped) reward achieved. Quantile regression can achieve this neatly by shaping the pile of earth according to an advantage estimation and the constraint of monotonicity (which is a core property of quantile functions).

More formally, we are interested in the effect of the quantile regression loss (1) on action probabilities when the error δ is between two samples from the same network, one representing the action that was played in the environment (which resulted from some sample quantile τ ) and one representing the quantile function value at some τ where the network tries to approximate the quantile function of the optimal action distribution. We focus in this analysis on the simple case of a single action dimension with scalar quantile target input and scalar quantile function output. The generalization to multidimensional action-/quantile-functions with independent action dimensions follows trivially.

First, let us denote the action taken in the environment by a τ := Qθ (τ, o) = Qθ (τ ), where we hide the dependence of Qθ (τ, o) on the observation o hereafter for notation simplicity. To train the quantile network we sample different τ and consider the loss

ρ τ (a τ -Qθ (τ )) = (τ -1 aτ -Qθ (τ )<0 ) • (τ -Qθ (τ )).

As we are interested in the effect of the loss on the probability of a τ consider a τ close to τ , i.e., τ = τ ± for some > 0 and < τ < 1 -. For τ = τthe loss reduces to

ρ τ (a τ -Qθ (τ -)) = (τ -) • (a τ -Qθ (τ -)) = (τ -) • a τ -Qθ (τ -) →0 ---→ dτ 1 δ Qθ (τ )

δτ where dτ 1 is an infinitely small but positive value Similarly we get for τ

= τ + ρ τ (a τ -Qθ (τ + )) = (τ + -1)•(a τ -Qθ (τ + )) = (1-τ -)• Qθ (τ + ) -a τ →0 ---→ dτ 2 δ Qθ (τ )

δτ with dτ 2 an infinitely small but positive value. Therefore, the quantile regression loss is positively correlated to the slope of Qθ at τ . The partial derivative of a quantile function Qθ (τ ) with respect to the quantile τ is however also known as the sparsity function (Tukey, 1965) or quantile-density function (Parzen, 1979) and has the interesting property (Jones, 1992) :

δ δτ Qθ (τ ) = 1 p Q θ ( Qθ (τ )) = 1 p Q θ (a τ )

where p Q θ is the probability density function of distribution Q θ . Hence, the quantile regression loss is inverse proportional to the probability of action a τ , which implies the following, given that Qθ is the quantile function of distribution Q θ :

1. Minimizing the quantile loss (1) for a given action a τ increases the likelihood of action a τ .

2. Maximizing the quantile loss for a given action a τ decreases the likelihood of action a τ . This leads us to the QRRL actor objective

min θ E aτ [E τ [ρ τ (a τ -Qθ (τ , o t )) • (R t,n -V ψ (o t )]] s.t. Qθ is monotonically increasing with τ where (R t,n -V ψ (o t )

) is an advantage estimation (Mnih et al., 2016) with R t,n = γ n V ψ (o t+n ) + t+n t =t γ t -t r t being the n-step estimate of the discounted future reward and V ψ (•) being the output of a critic network with parameters ψ, which is trained to approximate R t,n . Note that QRRL is a constraint optimization (since Qθ must be a proper quantile function, i.e., monotonically increasing).

In our experiments we found it however sufficient to add this constraint as an additional loss term

L mon (τ, τ , θ) = L κ Huber ( Qθ (τ ) -Qθ (τ )) • 1 τ <τ & Qθ (τ )> Qθ (τ ) or τ >τ & Qθ (τ )< Qθ (τ ) with L κ Huber (δ) = δ 2 |δ| ≤ κ |δ| -1

2 κ otherwise We weight this additional loss term with a constant Lagrange multiplier λ mon in the full loss term

L QRRL (a τ , o t , τ, τ , θ, ψ) = L a (a τ , o t , τ , θ) + λ c L c (o t , ψ) + λ mon L mon (τ, τ , θ) (2) with L a (a τ , o t , τ , θ) = ρ τ (a τ -Qθ (τ , o)) • (R t,n -V ψ (o t )) and L c (o t , ψ) = (R t,n -V ψ (o t )) 2 .

λ c is another constant Lagrange multiplier to weight the critic loss L c against the actor loss L a .

Although we focus in our experiments on an on-policy method with multiple actors, QRRL can easily be adapted to the off-policy setting.

4. STATE ALIGNED VECTOR REWARD AGENT:

In our State Aligned VEctor Reward (SAVER) agent, we use QRRL to train the agent network AN η through the PCPN. For this we feed the action probability distribution output µ, σ of the agent network to a pretrained PCPN and train on the QRRL loss (2) with respect to the agent network parameters η, where we take the actual position change ∆p introduced by a sampled action a ∼ N (µ, Iσ) as QRRL action target a τ = ∆p and compare it to K sampled position change estimations ∆ p(τ (i) ), τ (i) ∼ U([0, 1] d ) for i ∈ {1, ..., K}. Note that we pretrain the PCPN in our setup with random observations and action probability distributions and freeze the PCPN weights during agent training. Therefore, one could potentially train several agents to solve different tasks in the same environment using the same PCPN.

Since the output of the PCPN can be seen as state aligned action, training on vector rewards is straight forward. Instead of having a scalar critic estimating the value function V (•) ∈ R we estimate a value per action dimension, i.e., V (•) ∈ R d . Similarly, the vector rewards can be summed to a vector n-step discounted reward estimation R t,n = γ n V (o t+n ) + t+n t =t γ t -t r t , which leads us to a vector advantage (R t,n -V (o t )) at timestep t. Therefore we have an advantage estimation for each position change dimension individually and can apply loss (2) to each position change estimation dimension individually.

5. EXPERIMENTS AND RESULTS:

To test our ideas, we implement three simple experiments, two to test our approach in high dimensional metric spaces and one to show the applicability of our approach to a real world problem by modeling the robot arm shown in Figure 1(b) . For reproducibility and to stimulate further research in this area, our code is publicly available. 1 We compare our SAVER agent against an A2C agent, the synchronous variant of A3C (Mnih et al., 2016) . We choose this baseline since it is most similar to our SAVER implementation and we believe that SAVER can benefit in the future from state-ofthe-art additions to A3C as presented by Espeholt et al. (2018) . Here however, we want to focus on the benefit of using vector rewards instead of scalar rewards. We implement a simple feed forward In a first experiment, we let the agent move freely in all directions within a d-dimensional hypercube by choosing the action space to represent a set of (angle, step size) pairs for moving in the pairwise grouped environment dimensions. Note that this setup is similar to the depiction in Figure 1 (a), but the action dimensions are not aligned with the state-/position-change. The initial position of the agent is sampled uniformly at random from U([-1, 1] d ) and the goal of the agent is to reach the origin of the d-dimensional space. The agent's observation is its position and movement of the agent is stopped at the boundaries of the hypercube defined by

[-1, 1] d .

Step sizes are re-scaled and clipped to a maximal value of 0.1 and episodes are terminated after 1,000 steps if the agent didn't manage to get into close proximity of the goal beforehand. We pretrain the PCPN with 100,000 batches of 128 transitions each by randomly sampling

o ∼ U([-1, 1] d ), µ ∼ U([-1, 1] d ) and σ ∼ U([0, 1] d )

where actions between -1 and 1 correspond to step sizes between -0.1 and 0.1 and angles between 0 and π, respectively.2 We performed a small hyperparameter search for the learning rate in the 8-dimensional hypercube and settled for a learning rate of 0.01 (chosen from {0.01, 0.003, 0.001}) for SAVER and a learning rate of 0.0003 (chosen from {0.001, 0.0003, 0.0001}) for A2C. 3 We measure the performance of the agents by the mean length of the last 100 episodes. This mean episode length is plotted in Figure 3 over the course of training for hypercubes of different dimensionality.

Average and standard deviation of three training runs is shown. As can be seen from the plots, the higher the dimensionality d of the environment, the more apparent the advantage of training on vector rewards gets. SAVER trains faster in high dimensional cubes and is even able to find the goal in a 16-dimensional cube given the step limit of 1,000 steps.

As a second experiment, we keep the environment specifications the same but change the action representation. Here the agent's action consists of two parts: a softmax distribution from which the dimension to be manipulated is chosen (discrete action part) and a scalar Gaussian distribution defined by network outputs µ and σ from which the step size is sampled (continuous action part). Note that this action composition leads to a not continuous position change distribution when regarding the position change dimensions isolated, since in each dimension the probability of a change equal to 0 is more likely then position changes close to 0. We show with this experiment, that the quantile regression based PCPN can learn to implicitly model this complex position change distribution based on the composed action distribution input. For this we pretrain the PCPN with 10,000 batches of 128 transitions each by sampling softmax logits and µ uniformly at random from U([-1, 1]) and σ from U([0, 1]). In this experiment we searched for an appropriat learning rate in the 4 dimensional hypercube and settled for 0.01 for SAVER (chosen from {0.01, 0.003, 0.001}) and 0.0001 for A2C (chosen from {0.001, 0.0003, 0.0001}). The corresponding mean episode lengths over the course of agent training are plotted in Figure 4 . Again we find that the higher dimensional the hypercube is, the more advantageous it is to train with vector rewards.

In our last experiment, we model the 2-joint robot arm depicted in Figure 1 

6. RELATED WORK:

Using a vector of rewards instead of a scalar reward is most common in the literature on Multi-Objective Reinforcement Learning (MORL) (Liu et al., 2015) and Multi-Objective Sequential Decision Making (Roijers et al., 2013; Roijers & Whiteson, 2017) . However, since most of this literature focuses on classical reinforcement learning and conflicting objectives, the methods discussed either reduce the multiple objectives to a single objective or learn different policies for each objective additionally to a superpolicy deciding on which policy to use when. Recent work (Mossalam et al., 2016; Tajmajer, 2017; Nguyen, 2018) also applied these techniques to deep reinforcement learning agents. In contrast, we directly use the vector reward as more informative training signal for a neural network which implicitly learns to trade off different objectives.

Other adjacent areas are multi-agent and multi-task deep reinforcement learning (Arulkumaran et al., 2017) . While former considers multiple agents with potentially different observations, later addresses how a single agent best solves multiple tasks. In contrast, our work considers a single agent solving a single task by leveraging a multi-dimensional reward. Klinkhammer (2018) discusses multiple problem aligned rewards for better learning, while Brys et al. (2014) discuss the effect of correlated rewards on learning performance. Brys et al. (2014) also suggest multiple reward shapings of the same reward function for faster learning. Van Seijen et al. (2017) decompose the reward function into multiple rewards and train an architecture similar to ours with deep Q learning, assigning a Q-value output to each reward. While all three approaches come to the same conclusion as we do, i.e., increased training performance, they do require hand engineered reward functions, reward shapings and/or reward decompositions. In contrast, our approach is based on the fact that many state spaces are metric spaces and therefore allow for a straight forward vector reward interpretation. This makes our approach easier to apply to a larger set of tasks. While state proximity was already used by McCallum (1992) for faster backpropagation of rewards in tabular reinforcement learning, we are unaware of any deep learning algorithm capitalizing on state aligned vector rewards as we do.

Many recent approaches to deep reinforcement learning learn the environment dynamics to have a richer training signal (Dosovitskiy & Koltun, 2016) , imagine action consequences (Racanière et al., 2017) , improve exploration (Pathak et al., 2017) or dream up solutions (Ha & Schmidhuber, 2018 ). An interesting line of research (Zhang et al., 2017; Barreto et al., 2017; Ma et al., 2018; Barreto et al., 2018) in this direction uses successor features to share knowledge between tasks in the same environment. In contrast to most of these works, which often only predict one possible next state or successor feature, our PCPN incorporates the full probability distribution of possible state changes. While Ha & Schmidhuber (2018) also predict the full probability distribution of their next state representation by a Gaussian mixture model, our approach is more general in that it is also able to approximate non-Gaussian probability distributions. Dabney et al. (2018b) were the first to use quantile regression in connection with deep reinforcement learning. In their work, including their followup work (Dabney et al., 2018a) , they focused on approximating the full probability distribution of the value function. In contrast, with QRRL we explore possibilities of using quantile regression to approximate richer policies by not constraining the action distribution to an explicitly parameterized distribution. Ostrovski et al. (2018) showed that quantile networks can also be used for generative modeling. In general, we see quantile regression in combination with deep learning to have a lot of potential for future work.

7. CONCLUSION:

In this work we present the idea of state aligned vector rewards for faster reinforcement learning. While the idea is straight forward and simple, we are unaware of any work that addresses it so far. Additionally, we also present a new reinforcement learning technique based on quantile regression in this work which we term QRRL. QRRL allows for complex stochastic policies in continuous action spaces, not limiting agents to Gaussian actions. Combining both, we show that the agent network in our SAVER agent can be trained through a quantile network pretrained in the environment. We show that SAVER is capable of training orders of magnitudes faster in high dimensional metric spaces. While d-dimensional metric spaces are mainly mathematical constructs for d > 3, we see a lot of potential in SAVER to be applied to problems in mathematics and related fields, including the field of deep (reinforcement) learning itself. 

A ARCHITECTURE DETAILS:

As agent network, we used a simple layer-wise fully connected network with 3 hidden representations of size 256, 128 and 64, respectively. From the last hidden representation we map with a linear layer to the action mean µ, while the action standard deviation σ is initialized to e -1 • [1, ..., 1] T and kept independent of the observation input. For the A2C implementation, the last hidden representation of the agent network is also mapped through a linear layer to the scalar value estimate V (o). For the PCPN we first map from µ, σ and o to a hidden representation of size 256. From this hidden representation we map to d hidden representations, each of size 128, where d is the number of state dimensions. We then multiply each of these d hidden representation with a 128-dimensional cosine embedding of corresponding quantile target τ j with j ∈ {1, ..., d}. Each of this new hidden representations is then fed through a fully connected layer with 64 neurons before being linearly projected to a scalar value representing the position change estimation of dimension j for quantile τ j . We implement the QRRL critic as a separate network with 3 hidden representations of size 256, 128 and 64 and the same input as fed to the PCPN's first layer. Besides the action distribution representation which is ajusted corresponding to the experiment, we keep this architecture and all of the hyperparameters (as listed in Table 1 ) fixed for all experiments. The only hyperparameter we adjust is the learning rate, as described in the main text.

:

Our code can be found at: https://goo.gl/rMuadg (Anonymous Google Drive for double blind review)

Note that even though SAVER requires pretraining, pretraining is task independent and other tasks could be learned by reusing a trained PCPN. Therefore we do not include the pretraining in our plots. In our experiments however, we retrain the PCPN for each training run to show algorithmic stability.

We found in informal early experiments, that A2C requires a much smaller initial learning rate than the QRRL based SAVER for stable learning.

Figures and Tables:

Figure1: (a) An agent freely moving in a 2D world might try to reach a goal at position (0, 0) by taking action a. A sensible reward in this environment is the change in absolute distance to the goal. With a scalar reward this would be summarized as r = r x -r y , whereas a vector reward would keep the two reward dimensions distinguishable. (b) In most cases action and state space are however not aligned, therefore a mapping from action to state change must be learned.

Figure 2: An overview of the architecture setup. The agent and the position change prediction network (PCPN) are instantiated with neural networks and trained through quantile regression reinforcement learning and quantile regression, respectively.

Figure 3: Average number of steps needed to reach the goal over the course of training of an agent that chooses pairwise angle and step size for pairwise grouped dimensions. Results are shown for a 2, 4, 8 and 16 dimensional environment, corresponding plots are ordered from left to right. The x-axis shows the number of training steps in millions while the y-axis shows the average episode length over the last 100 episodes.Training length was fixed to 2,560,000 steps. Plotted is the average of 3 training runs with the shaded area indicating the standard deviation between runs. The plots suggest that the higher dimensional the environment, the more apparent the gain of training on vector rewards is.

Figure 4: Average number of steps needed to reach the goal over the course of training of an agent that can select the dimension it wants to move in and the amount by how much it wants to move. Results are shown for a 3, 4, 5 and 6 dimensional environment, corresponding plots are ordered from left to right. The x-axis shows the number of training steps in millions while the y-axis shows the average episode length over the last 100 episodes. Training length was fixed to 12,80,000 steps. Plotted is the average of 3 training runs with the shaded area indicating the standard deviation between runs. Again we see a benefit of training on vector rewards in high dimensional environments.

(b). Here an action sampled from the 2-dimensional Gaussian a ∼ N (µ, Iσ) is translated into an angle change between -0.02 • π and 0.02 • π of the joints. The agent's observation o consist here additionally to the hand position also of a goal position and the current angle of the joints. At the beginning of each episode, the initial joint angles and the goal position are randomly sampled from U([-π, π] 2 ) and U([-1, 1] 2 ), respectively. We choose the learning rate of SAVER as 0.003 and A2C as 0.0003 based on a search over {0.1, 0.01, 0.003, 0.001} and {0.01, 0.001, 0.0003, 0.0001}, respectively. We pretrain the PCPN with 10,000 batches of 128 transitions each by sampling µ ∼ U([-1, 1] d ) and σ ∼ U([0, 1] d ). The corresponding mean episode lengths over the course of training are plotted in Figure5. Even in this low dimensional problem, SAVER trains slightly faster than A2C due to the more informative vector rewards.

List of hyperparameters used in all experiments

Table content: <table><tr><td>Hyperparameter Value</td></tr><tr><td>Number of actors 16</td></tr><tr><td>Rollout length n 8</td></tr><tr><td>A2C entropy β 0.0</td></tr><tr><td>Critic λ c 0.5</td></tr><tr><td>Pretrain learning rate 0.1 • Learning rate</td></tr><tr><td>Monotonic λ mon 1.0</td></tr><tr><td>Discount factor γ 0.99</td></tr><tr><td>Quantile samples K 32</td></tr><tr><td>Optimizer RMSProp</td></tr><tr><td>Non-linearities ReLu</td></tr></table>

Peer review (iclr2019):
 The algorithm presented here aims at speeding up (reinforcement) learning in situations where the reward is “aligned” with the state space, i.e. can be decomposed into multiple components where each component depends only on one dimension of the state. Such a situation is found e.g. in tasks where the agent needs to reach a specific position in a 2D or 3D space, and the reward can be decomposed into a 2D / 3D reward vector associated to reaching the target coordinate along each axis. The algorithm consists in first pre-training (by collecting data obtained with random action distributions) a so-called “Position Change Prediction Network” (PCPN) to map an action distribution (e.g. mu and sigma if the agent’s policy is Gaussian) to the corresponding state change distribution along each relevant state dimension (for which there is an associated reward). The PCPN is an Implicit Quantile Network (Dabney et al, 2018), so as to model the state change distribution without assuming a specific parametric form. Once trained, the PCPN can be backpropagated through to train the agent distribution so that state changes that lead to better rewards (along each axis) become more frequent, while those leading to worse reward become less frequent (where “better” and “worse” are defined by the advantage, computed with a critic V). Experiments are performed on three toy tasks, and comparisons to A2C show faster learning for the proposed algorithm, especially in higher dimensional state spaces. The proposed approach combines two interesting ideas: (1) decomposing the reward into multiple components, and (2) aligning the agent’s action distribution with these components (through the PCPN). The second point is novel to the best of my knowledge. However, the first one is not since there exists a large body of literature on multi-objective reinforcement learning. The authors briefly mention some of it in the “Related Work” section, but unfortunately there is no comparison being made in experiments (the only algorithm being compared against is a “vanilla” A2C that uses a single scalar reward). Although previous work on multi-objective RL may not necessarily take advantage of the alignment between reward components and state dimensions, it is still relevant since the decomposition of the reward function by itself might be enough to speed-up learning, regardless of whether or not it is aligned with the state space. Experiments need to show that using the alignment as proposed here actually brings some benefits over existing multi-objective techniques. In addition to the lack of comparison to related multi-objective methods, the empirical evaluation is not very convincing since it is done on toy problems, with the most “realistic” one being a simple 2D task where the improvement brought by the proposed technique is not obvious (Fig. 5). Better results are obtained on high dimensional toy problems but (a) in these problems the relationship between actions and reward components is somewhat trivial, and (b) since the paper was originally motivated by spatial positioning tasks, it is not clear how often in practice one will need to go beyond 2 or 3 dimensions... Finally, since there is no theoretical analysis either (except for some intuition as to what happens in a limit case at the bottom of p. 4), and the optimized objective is essentially heuristically motivated, it is hard to tell how the algorithm would behave across a larger range of applications. I have two concerns in particular regarding the methodology: - Pre-training the PCPN assumes that randomly sampling action distributions will give us enough coverage of the state and action distribution spaces to generalize to those seen during training of the agent. This may not always be the case, requiring training the PCPN online: it would have been nice to verify whether this is feasible in practice. - I am afraid that forcing the agent to learn only from the PCPN gradient may prevent it from learning an optimal policy. As a dummy illustrative example, imagine an infinite 2D gridworld where the agent has two actions (action A = increase x coordinate by 1 only if y = 0, action B = increase both x and y by 1 only if y = 0), the reward is +1 when x increases and the agent starts at the origin. The reward is aligned with the x axis, and the optimal policy is to always take action A since as soon as action B is taken the agent can’t move anymore. With the PCPN output being the change along the x axis, its gradient wrt the agent action distribution is zero since both actions A and B increase x equally. Thus the agent will never be able to learn that A is to be preferred. In other words, it seems to me that the agent is limited to learn from changes in the state coordinates aligned with the reward, even if other coordinates are important for transition dynamics that may impact future rewards. Regarding the clarity of the paper, I found it good enough overall except that it is a bit difficult to follow the narrative. In particular Section 3 breaks the flow by considering a more general setting than the one introduced previously (the note about “action” and “policy” meaning something else is particularly confusing). Using Q to denote the action distribution is also a poor idea since Q is pretty much a “reserved” letter in RL. Overall the generic idea of “Quantile Regression RL” from Section 3 seems potentially interesting, and may be worth exploring on its own, but the current presentation (in the application to multi-objective state-aligned rewards) makes it more a distraction than an asset. Minor points: - The “biological insights” from the introduction are difficult to grasp since a reader versed in RL is unlikely to know what are “viral vector strategies” or “effector specific value estimations” - The definition of aligned spaces (“their dimensions correlate”) is pretty vague, a more precise mathematical definition would be helpful (even if the intuition is easy to grasp) - In the first one-line equation p. 4 there is a tau that should be a_tau - In the definition of R_t,n on p. 5, the sum should be up to t+n-1 (and same with the vector version at the bottom of the page) - I wonder if the loss L_mon is really needed, since the quantile regression loss should naturally lead to a monotonous function (assuming enough data and capacity). Are results significantly worse without it? - The definition of the Huber loss is incorrect - It is not explicitly stated that there is no backpropagation through V_psi in L_a - “QRRL can easily be adapted to the off-policy setting”: this is not obvious since you are using n-step returns bootstrapped with a critic V, which requires on-policy data (unless appropriate corrections are applied) - Experiments do not clearly define the rewards used by both algorithms - Figures would be easier to read if axes were labeled - I am surprised by the absence of entropy loss in A2C (beta = 0), was this hyper-parameter carefully tuned? - Having diagrams of the networks in Appendix A would be clearer than text explanations - The “cosine embedding” should be defined to avoid having to refer to Dabney et al - Note that the proposed model ends up having significantly more capacity than the single A2C network, which could raise concerns wrt the fairness of the comparison