Year: 2019

Abstract:
We consider reinforcement learning in input-driven environments, where an exogenous, stochastic input process affects the dynamics of the system. Input processes arise in many applications, including queuing systems, robotics control with disturbances, and object tracking. Since the state dynamics and rewards depend on the input process, the state alone provides limited information for the expected future returns. Therefore, policy gradient methods with standard state-dependent baselines suffer high variance during training. We derive a bias-free, input-dependent baseline to reduce this variance, and analytically show its benefits over state-dependent baselines. We then propose a meta-learning approach to overcome the complexity of learning a baseline that depends on a long sequence of inputs. Our experimental results show that across environments from queuing systems, computer networks, and MuJoCo robotic locomotion, input-dependent baselines consistently improve training stability and result in better eventual policies.

heading: 1 INTRODUCTION
text: Deep reinforcement learning (RL) has emerged as a powerful approach for sequential decision-making problems, achieving impressive results in domains such as game playing (Mnih et al., 2015; Silver et al., 2017) and robotics (Levine et al., 2016; Schulman et al., 2015a; Lillicrap et al., 2015). This paper concerns RL in input-driven environments. Informally, input-driven environments have dynamics that are partially dictated by an exogenous, stochastic input process. Queuing systems (Kleinrock, 1976; Kelly, 2011) are an example; their dynamics are governed by not only the decisions made within the system (e.g., scheduling, load balancing) but also the arrival process that brings work (e.g., jobs, customers, packets) into the system. Input-driven environments also arise naturally in many other domains: network control and optimization (Winstein & Balakrishnan, 2013; Mao et al., 2017), robotics control with stochastic disturbances (Pinto et al., 2017), locomotion in environments with complex terrains and obstacles (Heess et al., 2017), vehicular traffic control (Belletti et al., 2018; Wu et al., 2017), tracking moving targets, and more (see Figure 1).
We focus on model-free policy gradient RL algorithms (Williams, 1992; Mnih et al., 2016; Schulman et al., 2015a), which have been widely adopted and benchmarked for a variety of RL tasks (Duan et al., 2016; Wu & Tian, 2017). A key challenge for these methods is the high variance in the gradient estimates, as such variance increases sample complexity and can impede effective learning (Schulman et al., 2015b; Mnih et al., 2016). A standard approach to reduce variance is to subtract a “baseline” from the total reward (or “return”) to estimate the policy gradient (Weaver & Tao, 2001). The most common choice of a baseline is the value function — the expected return starting from the state.
Our main insight is that a state-dependent baseline — such as the value function — is a poor choice in input-driven environments, whose state dynamics and rewards are partially dictated by the input process. In such environments, comparing the return to the value function baseline may provide limited information about the quality of actions. The return obtained after taking a good action may be poor (lower than the baseline) if the input sequence following the action drives the system to unfavorable states; similarly, a bad action might end up with a high return with an advantageous input sequence. Intuitively, a good baseline for estimating the policy gradient should take the specific instance of the input process — the sequence of input values — into account. We call such a baseline an input-dependent baseline; it is a function of both the state and the entire future input sequence.
We formally define input-driven Markov decision processes, and we prove that an input-dependent baseline does not introduce bias in standard policy gradient algorithms such as Advantage Actor
Critic (A2C) (Mnih et al., 2016) and Trust Region Policy Optimization (TRPO) (Schulman et al., 2015a), provided that the input process is independent of the states and actions. We derive the optimal input-independent baseline and a simpler one to work with in practice; this takes the form of a conditional value function — the expected return given the state and the future input sequence.
Input-dependent baselines are harder to learn than their state-dependent counterparts; they are highdimensional functions of the sequence of input values. To learn input-dependent baselines efficiently, we propose a simple approach based on meta-learning (Finn et al., 2017; Vilalta & Drissi, 2002). The idea is to learn a “meta baseline” that can be specialized to a baseline for a specific input instantiation using a small number of training episodes with that input. This approach applies to applications in which an input sequence can be repeated during training, e.g., applications that use simulations or experiments with previously-collected input traces for training (McGough et al., 2017).
We compare our input-dependent baseline to the standard value function baseline for the five tasks illustrated in Figure 1. These tasks are derived from queuing systems (load balancing heterogeneous servers (Harchol-Balter & Vesilo, 2010)), computer networks (bitrate adaptation for video streaming (Mao et al., 2017)), and variants of standard continuous control RL benchmarks in the MuJoCo physics simulator (Todorov et al., 2012). We adapted three widely-used MuJoCo benchmarks (Duan et al., 2016; Clavera et al., 2018a; Heess et al., 2017) to add a stochastic input element that makes these tasks significantly more challenging. For example, we replaced the static target in a 7-DoF robotic arm target-reaching task with a randomly-moving target that the robot aims to track over time. Our results show that input-dependent baselines consistently provide improved training stability and better eventual policies. Input-dependent baselines are applicable to a variety of policy gradient methods, including A2C, TRPO, PPO, robust adversarial RL methods such as RARL (Pinto et al., 2017), and meta-policy optimization such as MB-MPO (Clavera et al., 2018b). Video demonstrations of our experiments are available at https://sites.google.com/view/input-dependent-baseline/.

heading: 2 PRELIMINARIES
text: Notation. We consider a discrete-time Markov decision process (MDP), defined by (S,A,P, ρ0, r, γ), where S ⊆ Rn is a set of n-dimensional states,A ⊆ Rm is a set ofm-dimensional actions, P : S ×A× S → [0, 1] is the state transition probability distribution, ρ0 : S → [0, 1] is the distribution over initial states, r : S ×A → R is the reward function, and γ ∈ (0, 1) is the discount factor. We denote a stochastic policy as π : S × A → [0, 1], which aims to optimize the expected return η(π) = Eτ [ ∑∞ t=0 γ
tr(st, at)], where τ = (s0, a0, ...) is the trajectory following s0 ∼ ρ0, at ∼ π(at|st), st+1 ∼ P(st+1|st, at). We use Vπ(st) = Eat,st+1,at+1,... [∑∞ l=0 γ lr(st+l, ar+l)|st ]
to define the value function, and Qπ(st, at) = Est+1,at+1,... [∑∞ l=0 γ lr(st+l, ar+l)|st, at ] to define the state-action value function. For any sequence (x0, x1, ...), we use x to denote the entire sequence and xi:j to denote (xi, xi+1, ..., xj).
Policy gradient methods. Policy gradient methods estimate the gradient of expected return with respect to the policy parameters (Sutton et al., 2000; Kakade, 2002; Gu et al., 2017). To train a policy
πθ parameterized by θ, the Policy Gradient Theorem (Sutton et al., 2000) states that
∇θη(πθ) = Es∼ρπ a∼πθ [∇θ log πθ(a|s)Qπθ (s, a)] , (1)
where ρπ(s) = ∑∞ t=0 [γ
t Pr(st = s)] denotes the discounted state visitation frequency. Practical algorithms often use the undiscounted state visitation frequency (i.e., γ = 1 in ρπ), which can make the estimation slightly biased (Thomas, 2014).
Estimating the policy gradient using Monte Carlo estimation for the Q function suffers from high variance (Mnih et al., 2016). To reduce variance, an appropriately chosen baseline b(st) can be subtracted from the Q-estimate without introducing bias (Greensmith et al., 2004). The policy gradient estimation with a baseline in Equation (1) becomes Eρπ,πθ [∇θ log πθ(a|s) (Qπθ (s, a)− b(s))]. While an optimal baseline exists (Greensmith et al., 2004; Wu et al., 2018), it is hard to estimate and often replaced by the value function b(st) = Vπ(st) (Sutton & Barto, 2017; Mnih et al., 2016).

heading: 3 MOTIVATING EXAMPLE
text: We use a simple load balancing example to illustrate the variance introduced by an exogenous input process. As shown in Figure 2a, jobs arrive over time and a load balancing agent sends them to one of two servers. The jobs arrive according to a Poisson process, and the job sizes follow a Pareto distribution. The two servers process jobs from their queues at identical rates. On each job arrival, the load balancer observes state st = (q1, q2), denoting the queue length at the two servers. It then takes an action at ∈ {1, 2}, sending the job to one of the servers. The goal of the load balancer is to minimize the average job completion time. The reward corresponding to this goal is rt = −τ × j, where τ is the time elapsed since the last action and j is total number of enqueued jobs.
In this example, the optimal policy is to send the job to the server with the shortest queue (Daley, 1987). However, we find that a standard policy gradient algorithm, A2C (Mnih et al., 2016), trained using a value function baseline struggles to learn this policy. The reason is that the stochastic sequence of job arrivals creates huge variance in the reward signal, making it difficult to distinguish between good and bad actions. Consider, for example, an action at the state shown in Figure 2a. If the arrival sequence following this action consists of a burst of large jobs (e.g., input sequence 1 in Figure 2a), the queues will build up, and the return will be poor compared to the value function baseline (average return from the state). On the other hand, a light stream of jobs (e.g., input sequence 2 in Figure 2a) will lead to short queues and a better-than-average return. Importantly, this difference in return has little to do with the action; it is a consequence of the random job arrival process.
We train two A2C agents (Mnih et al., 2016), one with the standard value function baseline and the other with an input-dependent baseline tailored for each specific instantiation of the job arrival process (details of this baseline in §4). Since the the input-dependent baseline takes each input sequence into account explicitly, it reduces the variance of the policy gradient estimation much more effectively (Figure 2b, left). As a result, even in this simple example, only the policy learned with the input-dependent baseline comes close to the optimal (Figure 2b, right). Figure 2c visualizes the policies learned using the two baselines. The optimal policy (pick-shortest-queue) corresponds to a clear divide between the chosen servers at the diagonal.
In fact, the variance of the standard baseline can be arbitrarily worse than an input-dependent baseline: we refer the reader to Appendix A for an analytical example on a 1D grid world.

heading: 4 REDUCING VARIANCE FOR INPUT-DRIVEN MDPS
text: We now formally define input-driven MDPs and derive variance-reducing baselines for policy gradient methods in environments with input processes. Definition 1. An input-driven MDP is defined by (S,A,Z,Ps,Pz, ρs0, ρz0, r, γ), where Z ⊆ Rk is a set of k-dimensional input values, Ps(st+1|st, at, zt) is the transition kernel of the states, Pz(zt+1|z0:t) is the transition kernel of the input process, ρz0(z0) is the distribution of the initial input, r(st, at, zt) is the reward function, and S, A, ρs0, γ follow the standard definition in §2.
An input-driven MDP adds an input process, z = (z0, z1, · · · ), to a standard MDP. In this setting, the next state st+1 depends on (st, at, zt). We seek to learn policies that maximize cumulative expected rewards. We focus on two cases, corresponding to the graphical models shown in Figure 3: Case 1: zt is a Markov process, and ωt = (st, zt) is observed at time t. The action at can hence depend on both st and zt. Case 2: zt is a general process (not necessarily Markov), and ωt = st is observed at time t. The action at hence depends only on st.
In Appendix B, we prove that case 1 corresponds to a fully-observable MDP. This is evident from the graphical model in Figure 3a by considering ωt = (st, zt) to be the ‘state’ of the MDP at time t. Case 2, on the other hand, corresponds to a partially-observed MDP (POMDP) if we define the state to contain both st and z0:t, but leave z0:t unobserved at time t (see Appendix B for details).

heading: 4.1 VARIANCE REDUCTION
text: In input-driven MDPs, the standard input-agnostic baseline is ineffective at reducing variance, as shown by our motivating example (§3). We propose to use an input-dependent baseline of the form b(ωt, zt:∞)— a function of both the observation at time t and the input sequence from t onwards. An input-dependent baseline uses information that is not available to the policy. Specifically, the input sequence zt:∞ cannot be used when taking an action at time t, because zt+1:∞ has not yet occurred at time t. However, in many applications, the input sequence is known at training time. In some cases, we know the entire input sequence upfront, e.g., when training in a simulator. In other situations, we can record the input sequence on the fly during training. Then, after a training episode, we can use the recorded values, including those that occurred after time t, to compute the baseline for each step t.
We now analyze input-dependent baselines. Our main result is that input-dependent baselines are bias-free. We also derive the optimal input-dependent baseline for variance reduction. All the results hold for both cases in Figure 3. We first state two useful lemmas required for our analysis. The first lemma shows that under the input-driven MDP definition, the input sequence zt:∞ is conditionally independent of the action at given the observation ωt, while the second lemma states the policy gradient theorem for input-driven MDPs. Lemma 1. Pr(zt:∞, at|ωt) = Pr(zt:∞|ωt)πθ(at|ωt), i.e., zt:∞ − ωt − at forms a Markov chain. Proof. See Appendix C. Lemma 2. For an input-driven MDP, the policy gradient theorem can be rewritten as
∇θη (πθ) = E(ω,z)∼ρπ a∼πθ
[ ∇θ log πθ(a|ω)Q(ω, a,z) ] , (2)
where ρπ(ω,z) = ∑∞ t=0 [γ
t Pr(ωt = ω, zt:∞ = z)] denotes the discounted visitation frequency of the observation ω and input sequence z, and Q(ω, a, z) = E [∑∞
l=0 γ lrt+l ∣∣ ωt = ω, at = a, zt:∞ = z]. Proof. See Appendix D.
Equation (2) generalizes the standard Policy Gradient Theorem in Equation (1). ρπ(ω,z) can be thought of as a joint distribution over observations and input sequences. Q(ω, a, z) is a “state-actioninput” value function, i.e., the expected return when taking action a after observing ω, with input sequence z from that step onwards. The key ingredient in the proof of Lemma 2 is the conditional independence of the input process zt:∞ and the action at given the observation ωt (Lemma 1).
Theorem 1. An input-dependent baseline does not bias the policy gradient.
Proof. Using Lemma 2, we need to show: E(ω,z)∼ρπ,a∼πθ [∇θ log πθ(a|ω)b(ω,z)] = 0. We have:
E(ω,z)∼ρπ a∼πθ [∇θ log πθ(a|ω)b(ω,z)] = ∑ ω ∑ z ∑ a ρπ(ω,z)πθ(a|ω)∇θ log πθ(a|ω)b(ω,z)
= ∑ ω ∑ z ρπ(ω,z)b(ω,z) ∑ a πθ(a|ω)∇θ log πθ(a|ω). (3)
Since ∑ a πθ(a|ω)∇θ log πθ(a|ω) = ∑ a∇θπθ(a|ω) = ∇θ ∑ a πθ(a|ω) = 0, the theorem follows.
Input-dependent baselines are also bias-free for policy optimization methods such as TRPO (Schulman et al., 2015a), as we show in Appendix F. Next, we derive the optimal input-dependent baseline for variance reduction. As the gradient estimates are vectors, we use the trace of the covariance matrix as the minimization objective (Greensmith et al., 2004).
Theorem 2. The input-dependent baseline that minimizes variance in policy gradient is given by
b∗(ω,z) = Ea∼πθ
[ ∇θ log πθ(a|ω)T∇θ log πθ(a|ω)Q(ω, a, z) ] Ea∼πθ [∇θ log πθ(a|ω)T∇θ log πθ(a|ω)] . (4)
Proof. See Appendix E.
Operationally, for observation ωt at each step t, the input-dependent baseline takes the form b(ωt, zt:∞). In practice, we use a simpler alternative to Equation (4): b(ωt, zt:∞) = Eat∼πθ [Q(ωt, at, zt:∞)]. This can be thought of as a value function V (ωt, zt:∞) that provides the expected return given observation ωt and input sequence zt:∞ from that step onwards. We discuss how to estimate input-dependent baselines efficiently in §5.
Remark. Input-dependent baselines are generally applicable to reducing variance for policy gradient methods in input-driven environments. In this paper, we apply input-dependent baselines to A2C (§6.2), TRPO (§6.1) and PPO (Appendix L). Our technique is complementary and orthogonal to adversarial RL (e.g., RARL (Pinto et al., 2017)) and meta-policy adaptation (e.g., MB-MPO (Clavera et al., 2018b)) for environments with external disturbances. Adversarial RL improves policy robustness by co-training an “adversary” to generate a worst-case disturbance process. Meta-policy optimization aims for fast policy adaptation to handle model discrepancy between training and testing. By contrast, input-dependent baselines improve policy optimization itself in the presence of stochastic input processes. Our work primarily focuses on learning a single policy in input-driven environments, without policy adaptation. However, input-dependent baselines can be used as a general method to improve the policy optimization step in adversarial RL and meta-policy adaptation methods. For example, in Appendix M, we empirically show that if an adversary generates high-variance noise, RARL with a standard state-based baseline cannot train good controllers, but the input-dependent baseline helps improve the policy’s performance. Similarly, input-dependent baselines can improve meta-policy optimization in environments with stochastic disturbances, as we show in Appendix N.

heading: 5 LEARNING INPUT-DEPENDENT BASELINES EFFICIENTLY
text: Input-dependent baselines are functions of the sequence of input values. A natural approach to train such baselines is to use models that operate on sequences (e.g., LSTMs (Gers et al., 1999)). However, learning a sequential mapping in a high-dimensional space can be expensive (Bahdanau et al., 2014). We considered an LSTM approach, but ruled it out when initial experiments showed that it fails to provide significant policy improvement over the standard baseline in our environments (Appendix G).
Fortunately, we can learn the baseline much more efficiently in applications where we can repeat the same input sequence multiple times during training. Input-repeatability is feasible in many applications: it is straightforward when using simulators for training, and also feasible when training a real system with previously-collected input traces outside simulation. For example, training a robot in the presence of exogenous forces might apply a set of time-series traces of these forces repeatedly to the physical robot. We now present two approaches that exploit input-repeatability to learn input-dependent baselines efficiently.
Multi-value-network approach. A straightforward way to learn b(ωt, zt:∞) for different input instantiations z is to train one value network to each particular instantiation of the input process. Specifically, in the training process, we first generate N input sequences {z1, z2, · · · , zN} and restrict training only to those N sequences. To learn a separate baseline function for each input sequence, we use N value networks with independent parameters θV1 , θV2 , · · · , θVN , and single policy network with parameter θ. During training, we randomly sample an input sequence zi, execute a rollout based on zi with the current policy πθ, and use the (state, action, reward) data to train the value network parameter θVi and the policy network parameter θ (details in Appendix I).
Meta-learning approach. The multi-value-network approach does not scale if the task requires training over a large number of input instantiations to generalize. The number of inputs needed is environment-specific, and can depend on a variety of factors, such as the time horizon of the problem, the distribution of the input process, the relative magnitude of the variance due to the input process compared to other sources of randomness (e.g., actions). Ideally, we would like an approach that enables learning across many different input sequences. We present a method based on meta-learning to train with an unbounded number of input sequences. The idea is to use all (potentially infinitely many) input sequences to learn a “meta value network” model. Then, for each specific input sequence, we first customize the meta value network using a few example rollouts with that input sequence. We then compute the actual baseline values for training the policy network parameters, using the customized value network for the specific input sequence. Our implementation uses Model-Agnostic Meta-Learning (MAML) (Finn et al., 2017).
Algorithm 1 Training a meta input-dependent baseline for policy-based methods. Require: α, β: meta value network step size hyperparameters
1: Initialize policy network parameters θ and meta-value-network parameters θV 2: while not done do 3: Generate a new input sequence z 4: Sample k rollouts T1, T2, ..., Tk using policy πθ and input sequence z 5: Adapt θV with the first k/2 rollouts: θ1V = θV − α∇θV LT1:k/2 [VθV ] 6: Estimate baseline value Vθ1V (ωt) for st ∼ Tk/2:k using adapted θ 1 V 7: Adapt θV with the second k/2 rollouts: θ2V = θV − α∇θV LTk/2:k [VθV ] 8: Estimate baseline value Vθ2V (ωt) for st ∼ T1:k/2 using adapted θ 2 V
9: Update policy with Equation (2) using the values from line (6) and (8) as baseline 10: Update meta value network: θV ← θV − β∇θV Lk/2:k [ Vθ1V ] − β∇θV L1:k/2 [ Vθ2V ] 11: end while
The pseudocode in Algorithm 1 depicts the training algorithm. We follow the notation of MAML, denoting the loss in the value function VθV (·) on a rollout T as LT [VθV ] = ∑ ωt,rt∼T ‖VθV (ωt)−∑T
t′=t γ t′−trt‖2. We perform rollouts k times with the same input sequence z (lines 3 and 4); we use the first k/2 rollouts to customize the meta value network for this instantiation of z (line 5), and then apply the customized value network on the states of the other k/2 rollouts to compute the baseline for those rollouts (line 6); similarly, we swap the two groups of rollouts and repeat the same process (lines 7 and 8). We use different rollouts to adapt the meta value network and compute the baseline to avoid introducing extra bias to the baseline. Finally, we use the baseline values computed for each rollout to update the policy network parameters (line 9), and we apply the MAML (Finn et al., 2017) gradient step to update the meta value network model (line 10).

heading: 6 EXPERIMENTS
text: Our experiments demonstrate that input-dependent baselines provide consistent performance gains across multiple continuous-action MuJoCo simulated robotic locomotions and discrete-action environments in queuing systems and network control. We conduct experiments for both policy gradient methods and policy optimization methods (see Appendix K for details). The videos for our experiments are available at https://sites.google.com/view/input-dependent-baseline/.

heading: 6.1 SIMULATED ROBOTIC LOCOMOTION
text: We use the MuJoCo physics engine (Todorov et al., 2012) in OpenAI Gym (Brockman et al., 2016) to evaluate input-dependent baselines for robotic control tasks with external disturbance. We extend the standard Walker2d, HalfCheetah and 7-DoF robotic arm environments, adding a different external input to each (Figure 1).
Walker2d with random wind (Figure 1c). We train a 2D walker with varying wind, which randomly drags the walker backward or forward with different force at each step. The wind vector changes randomly, i.e., the wind forms a random input process. We add a force sensor to the state to enable the agent to quickly adapt. The goal is for the walker to walk forward while keeping balance.
HalfCheetah on floating tiles with random buoyancy (Figure 1d). A half-cheetah runs over a series of tiles floating on water (Clavera et al., 2018a). Each tile has different damping and friction properties, which moves the half-cheetah up and down and changes its dynamics. This random buoyancy is the external input process; the cheetah needs to learn running forward over varying tiles.
7-DoF arm tracking moving target (Figure 1e). We train a simulated robot arm to track a randomly moving target (a red ball). The robotic arm has seven degrees of freedom and the target is doing a random walk, which forms the external input process. The reward is the negative squared distance between the robot hand (blue square) and the target.
The Walker2d and 7-DoF arm environments correspond to the fully observable MDP case in Figure 3, i.e. the agent observes the input zt at time t. The HalfCheetah environment is a POMDP, as the agent does not observe the buoyancy of the tiles. In Appendix H, we show results for the POMDP version of the Walker2d environment.
Results. We build 10-value networks and a meta-baseline using MAML, both on top of the OpenAI’s TRPO implementation (Dhariwal et al., 2017). Figure 4 shows the performance comparison among different baselines with 100 unseen testing input sequences at each training checkpoint. These learning curves show that TRPO with a state-dependent baseline performs worst in all environments. With the input-dependent baseline, by contrast, performance in unseen testing environments improves by up to 3×, as the agent learns a policy robust against disturbances. For example, it learns to lean into headwind and quickly place its leg forward to counter the headwind; it learns to apply different force on tiles with different buoyancy to avoid falling over; and it learns to co-adjust multiple joints to keep track of the moving object. The meta-baseline eventually outperforms 10-value networks as it effectively learns from a large number of input processes and hence generalizes better.
The input-dependent baseline technique applies generally on top of policy optimization methods. In Appendix L, we show a similar comparison with PPO (Schulman et al., 2017). Also, in Appendix M we show that adversarial RL (e.g., RARL (Pinto et al., 2017)) alone is not adequate to solve the high variance problem, and the input-dependent baseline helps improve the policy performance (Figure 9).

heading: 6.2 DISCRETE-ACTION ENVIRONMENTS
text: Our discrete-action environments arise from widely-studied problems in computer systems research: load balancing and bitrate adaptation.1 As these problems often lack closed-form optimal solutions (Grandl et al., 2016; Yin et al., 2015), hand-tuned heuristics abound. Recent work suggests
1We considered Atari games often used as benchmark discrete-action RL environments (Mnih et al., 2015). However, Atari games lack an exogenous input process: a random seed perturbs the games’ initial state, but it does not affect the environmental changes (e.g., in “Seaquest”, the ships always come in a fixed pattern).
that model-free reinforcement learning can achieve better performance than such human-engineered heuristics (Mao et al., 2016; Evans & Gao, 2016; Mao et al., 2017; Mirhoseini et al., 2017). We consider a load balancing environment (similar to the example in §3) and a bitrate adaptation environment in video streaming (Yin et al., 2015). The detailed setup of these environments is in Appendix J.
Results. We extend OpenAI’s A2C implementation (Dhariwal et al., 2017) for our baselines. The learning curves in Figure 5 illustrate that directly applying A2C with a standard value network as the baseline results in unstable test reward and underperforms the traditional heuristic in both environments. Our input-dependent baselines reduce the variance and improve test reward by 25–33%, outperforming the heuristic. The meta-baseline performs the best in all environments.

heading: 7 RELATED WORK
text: Policy gradient methods compute unbiased gradient estimates, but can experience a large variance (Sutton & Barto, 2017; Weaver & Tao, 2001). Reducing variance for policy-based methods using a baseline has been shown to be effective (Williams, 1992; Sutton & Barto, 2017; Weaver & Tao, 2001; Greensmith et al., 2004; Mnih et al., 2016). Much of this work focuses on variance reduction in a general MDP setting, rather than variance reduction for MDPs with specific stochastic structures. Wu et al. (2018)’s techniques for MDPs with multi-variate independent actions are closest to our work. Their state-action-dependent baseline improves training efficiency and model performance on high-dimensional control tasks by explicitly factoring out, for each action, the effect due to other actions. By contrast, our work exploits the structure of state transitions instead of stochastic policy.
Recent work has also investigated the bias-variance tradeoff in policy gradient methods. Schulman et al. (2015b) replace the Monte Carlo return with a λ-weighted return estimation (similar to TD(λ) with value function bootstrap (Tesauro, 1995)), improving performance in high-dimensional control tasks. Other recent approaches use more general control variates to construct variants of policy gradient algorithms. Tucker et al. (2018) compare the recent work, both analytically on a linearquadratic-Gaussian task and empirically on complex robotic control tasks. Analysis of control variates for policy gradient methods is a well-studied topic, and extending such analyses (e.g., Greensmith et al. (2004)) to the input-driven MDP setting could be interesting future work.
In other contexts, prior work has proposed new RL training methodologies for environments with disturbances. Clavera et al. (2018b) adapts the policy to different pattern of disturbance by training the RL agent using meta-learning. RARL (Pinto et al., 2017) improves policy robustness by co-training an adversary to generate a worst-case noise process. Our work is orthogonal and complementary to these work, as we seek to improve policy optimization itself in the presence of inputs like disturbances.

heading: 8 CONCLUSION
text: We introduced input-driven Markov Decision Processes in which stochastic input processes influence state dynamics and rewards. In this setting, we demonstrated that an input-dependent baseline can significantly reduce variance for policy gradient methods, improving training stability and the quality of learned policies. Our work provides an important ingredient for using RL successfully in a variety of domains, including queuing networks and computer systems, where an input workload is a fundamental aspect of the system, as well as domains where the input process is more implicit, like robotics control with disturbances or random obstacles.
We showed that meta-learning provides an efficient way to learn input-dependent baselines for applications where input sequences can be repeated during training. Investigating efficient architectures for input-dependent baselines for cases where the input process cannot be repeated in training is an interesting direction for future work.
Acknowledgements. We thank Ignasi Clavera for sharing the HalfCheetah environment, Jonas Rothfuss for the comments on meta-policy optimization and the anonymous ICLR reviewers for their feedback. This work was funded in part by NSF grants CNS-1751009, CNS-1617702, a Google Faculty Research Award, an AWS Machine Learning Research Award, a Cisco Research Center Award, an Alfred P. Sloan Research Fellowship and the sponsors of MIT Data Systems and AI Lab.

heading: B MARKOV PROPERTIES OF INPUT-DRIVEN DECISION PROCESSES
text: Proposition 1. An input-driven decision process satisfying the conditions of case 1 in Figure 3 is a fully observable MDP, with state s̃t := (st, zt), and action ãt := at.
Proof.
Pr(s̃t+1|s̃0:t, ã0:t) = Pr(st+1, zt+1|s0:t, z0:t, a0:t) = Pr(st+1, zt+1|st, zt, at) (by definition of case 1 in Figure 3a) = Pr(s̃t+1|s̃t, ãt).
Proposition 2. An input-driven decision process satisfying the conditions of case 2 in Figure 3, with state s̃t := (st, z0:t) and action ãt := at is a fully observable MDP. If only ωt = st is observed at time t, it is a partially observable MDP (POMDP).
Proof.
Pr(s̃t+1|s̃0:t, ã0:t) = Pr(st+1, z0:t+1|s0:t, z0:t, a0:t) = Pr(st+1|s0:t, z0:t+1, a0:t) Pr(z0:t+1|s0:t, z0:t, a0:t) = Pr(st+1|st, z0:t+1, at) Pr(z0:t+1|st, z0:t, at) (by definition of case 2 in Figure 3b) = Pr(st+1, z0:t+1|st, z0:t, at) = Pr(s̃t+1|s̃t, ãt).
Therefore, (s̃t, ãt) is a fully observable MDP. If only ωt = st is observed, the decision process is a POMDP, since the z0:t component of the state is not observed.

heading: C PROOF OF LEMMA 1
text: Proof. From the definition of an input-driven MDP (Definition 1), we have
Pr(z0:∞, ωt, at) =Pr(z0:t, ωt, at) Pr(zt+1:∞|z0:t, ωt, at) =Pr(z0:t, ωt) Pr(at|z0:t, ωt) Pr(zt+1:∞|z0:t) =Pr(z0:t, ωt)πθ(at|ωt) Pr(zt+1:∞|z0:t) =Pr(z0:∞, ωt)πθ(at|ωt). (9)
Notice that Pr(at|z0:t, ωt) = πθ(at|ωt) in both the MDP and POMDP cases in Figure 3. By marginalizing over z0:t−1 on both sides, we obtain the result:
Pr(zt:∞, ωt, at) = Pr(zt:∞, ωt)πθ(at|ωt). (10)

heading: D PROOF OF LEMMA 2
text: Proof. Expanding the Policy Gradient Theorem (Sutton & Barto, 2017), we have
∇θη(πθ) =E  ∞∑ t=0 ∇θ log πθ(at|ωt) ∑ t′≥t γt ′ rt′  =
∞∑ t=0 E ∇θ log πθ(at|ωt)∑ t′≥t γt ′ rt′  =
∞∑ t=0 ∑ ω,a,z Pr(ωt = ω, at = a, zt:∞ = z)∇θ log πθ(a|ω)E ∑ t′≥t γt ′ rt′ |ωt = ω, at = a, zt:∞ = z  =
∞∑ t=0 ∑ ω,a,z Pr(ωt = ω, zt:∞ = z)πθ(a|ω)∇θ log πθ(a|ω)E ∑ t′≥t γt ′ rt′ |ωt = ω, at = a, zt:∞ = z  , (11)
where the last step uses Lemma 1. Using the definition of Q(ω, a,z), we obtain:
∇θη(πθ) = ∞∑ t=0 [∑ ω,a,z Pr(ωt = ω, zt:∞ = z)πθ(a|ω)∇θ log πθ(a|ω)γtQ(ω, a, z) ]
= ∑ ω,a,z
[ πθ(a|ω)∇θ log πθ(a|ω)Q(ω, a,z) [ ∞∑ t=0 γt Pr(ωt = ω, zt:∞ = z) ]] = ∑ ω,a,z πθ(a|ω)∇θ log πθ(a|ω)Q(ω, a, z)ρπ(ω,z)
=E(ω,z)∼ρπ a∼πθ
[ ∇θ log πθ(a|ω)Q(ω, a,z) ] . (12)

heading: E PROOF OF THEOREM 2
text: Proof. Let G(ω, a) denote ∇θ log πθ(a|ω)T∇θ log πθ(a|ω). For any input-dependent baseline b(ω,z), the variance of the policy gradient estimate is given by
E(ω,z)∼ρπ a∼πθ
[ ‖∇θ log πθ(a|ω) [Q(ω, a,z)− b(ω,z)]− Eρπ,πθ [∇θ log πθ(a|ω) [Q(ω, a, z)− b(ω,z)]]‖ 2 2 ] = Eρπ,πθ [ G(ω, a) [ Q(ω, a, z)− b(ω,z) ]2]− ∥∥∥Eρπ,πθ[∇θ log πθ(a|ω)[Q(ω, a, z)− b(ω,z)]]∥∥∥2 2
= Eρπ,πθ [ G(ω, a) [ Q(ω, a,z)− b(ω,z) ]2]− ∥∥∥Eρπ,πθ[∇θ log πθ(a|ω)Q(ω, a, z)]∥∥∥2 2 (due to Theorem 1)
= Eρπ,πθ [ G(ω, a)Q(ω, a, z)2 ] − ‖Eρπ,πθ [∇θ log πθ(a|ω)Q(ω, a, z)]‖ 2 2
+ Eρπ [ Ea∼πθ [ G(ω, a) ∣∣ z, ω] b(ω,z)2 − 2Ea∼πθ [G(ω, a)Q(ω, a, z) ∣∣ ω,z] b(ω,z)]. Notice that the baseline is only involved in the last term in a quadratic form, where the second order term is positive. To minimize the variance, we set baseline to the minimizer of the quadratic equation, i.e., 2Ea∼πθ [ G(ω, a)
∣∣ ω,z] b(ω,z)− 2Ea∼πθ [G(ω, a)Q(ω, a, z) ∣∣ ω,z] = 0 and hence the result follows.
F INPUT-DEPENDENT BASELINE FOR TRPO
We show that the input-dependent baselines are bias-free for Trust Region Policy Optimization (TRPO) (Schulman et al., 2015a).
Preliminaries. Stochastic gradient descent using Equation (1) does not guarantee consistent policy improvement in complex control problems. TRPO is an alternative approach that offers monotonic policy improvements, and derives a practical algorithm with better sample efficiency and performance. TRPO maximizes a surrogate objective, subject to a KL divergence constraint:
maximize θ Es∼ρπold a∼πold [ πθ(a|s) πold(a|s) Qπold(s, a) ] (13)
subject to Es∼ρπold [DKL (πold(·|s)||πθ(·|s))] ≤ δ, (14)
in which δ serves as a step size for policy update. Using a baseline in the TRPO objective, i.e. replacing Qπold(s, a) with Qπold(s, a)− b(s), empirically improves policy performance (Schulman et al., 2015b). Similar to Theorem 2, we generalize TRPO to input-driven environments, with ρπ(ω,z) =∑∞ t=0 [γ
t Pr(ωt = ω, zt:∞ = z)] denoting the discounted visitation frequency of the observation ω and input sequence z, and Q(ω, a, z) = E [∑∞ l=0 γ lrt+l ∣∣ ωt = ω, at = a, zt:∞ = z]. The TRPO objective becomes E(ω,z)∼ρold,a∼πold [Qπold(ω, a, z)πθ(a|ω)/πold(a|ω)], and the constraint is E(ω,z)∼ρπold [DKL (πold(·|s)||πθ(·|s))] ≤ δ. Theorem 3. An input-dependent baseline does not change the optimal solution of the optimization problem in TRPO, that is argmaxθE(ω,z)∼ρold,a∼πold [ πθ(a|ω) πold(a|ω)Qπold(ω, a, z) ] =
argmaxθE(ω,z)∼ρold,a∼πold [ πθ(a|ω) πold(a|ω) (Qπold(ω, a,z)− b(ω,z)) ] .
Proof.
E(ω,z)∼ρold,a∼πold [ πθ(a|ω) πold(a|ω) b(ω,z) ] = ∑ ω ∑ z ρold(ω,z) ∑ a πold(a|ω) [ πθ(a|ω) πold(a|ω) b(ω,z) ] = ∑ ω ∑ z ρold(ω,z) ∑ a πθ(a|ω)b(ω,z)
= ∑ ω ∑ z ρold(ω,z)b(ω,z),
which is independent of θ. Therefore, b(ω,z) does not change the optimal solution to the optimization problem.
G INPUT-DEPENDENT BASELINE WITH LSTM
The input-dependent baseline is a function of both the state and the entire future input sequence. A natural approach to approximate such baselines is to use neural models that operate on sequences (e.g., LSTMs (Gers et al., 1999)). However, learning a sequential mapping in a high-dimensional space can be expensive (Bahdanau et al., 2014). For example, consider the LSTM input-dependent baseline with A2C on the load balancing environment (Figure 1a; §6.2) and TRPO on the Walker2d with wind environment (Figure 1c; §6.1). As shown in Figure 6, the LSTM input-dependent baseline (green) does not significantly improve the policy performance over a standard state-only baseline (red) in these environments. By contrast, the MAML based input-dependent baseline (§5) reduces the variance in policy gradient estimation much more effectively and achieves a consistently better policy performance.
H ADDITIONAL POMDP EXPERIMENT
Recall that the HalfCheetah on floating tiles environment (Figure 1d) is a POMDP, since the agent does not observe the buoyancy of the tiles. Figure 4 (middle) shows that the input-dependent baseline significantly improves the TRPO performance for this POMDP environment. We also created a POMDP version of the Walker2d with wind environment (Figure 1c), where the force of the wind is removed from the observation provided to the agent. The results of repeating the same training are shown in Figure 7. We make two key observations: (1) Compared with the MDP case in Figure 4 (left), the performance drops slightly overall. This is expected because the agent can react more agilely if it directly observes the current wind in the MDP case. (2) Input-dependent baselines reduce variance and improve the policy performance, with the MAML-based approach achieving the best performance, similar to the MDP case.

heading: I PSEUDOCODE FOR TRAINING MULTI-VALUE BASELINES
text: In §5, we explained the idea of efficiently computing input-dependent baselines (§4.1) using multiple value networks on a fixed set of input sequences. Algorithm 2 depicts the details of this approach.
Algorithm 2 Training multi-value baselines for policy-based methods. Require: pregenerated input seuqnces {z1, z2, · · · , zN}, step sizes α, β
1: Initialize value network parameters θV1 , θV1 , · · · , θVN and policy parameters θ 2: while not done do 3: Sample a input sequence zi 4: Sample k rollouts T1, T2, ..., Tk using policy πθ and input sequence zi 5: Update policy with Equation (2) using baseline estimated with θVi 6: Update i-th value network parameters: θVi ← θVi − β∇θViL1:k [ VθVi
] 7: end while

heading: J SETUP FOR DISCRETE-ACTION ENVIRONMENTS
text: Load balancing across servers (Figure 1a). In this environment, an RL agent balances jobs over k servers to minimize the average job completion time. Similar to §3, the job sizes follow a Pareto distribution (scale xm = 100, shape α = 1.5), and jobs arrive in a Poisson process (λ = 55). We run over 10 simulated servers with different processing rates, ranging linearly from 0.15 to 1.05. In this setting, the load of the system is at 90% (i.e., on average, 90% of the queues are non-empty). In each episode, we generate 500 jobs as the exogenous input process. The problem of minimizing average job completion time on servers with heterogeneous processing rates does not have a closed-form solution (Harchol-Balter & Vesilo, 2010); the most widely-used heuristic is to join the shortest queue (Daley, 1987). However, understanding the workload pattern can give a better policy; for example, we can reserve some servers for small jobs. In this environment, the observed state is a vector of (j, q1, q2, ..., qk), where j is the size of the incoming job, qi is the amount of work currently in each queue. The action a ∈ {1, 2, ..., k} schedules the incoming job to a specific queue. The reward is the number of active jobs times the negated time elapsed since the last action.
Bitrate adaptation for video streaming (Figure 1b). Streaming video over variable-bandwidth connections requires the client to adapt the video bitrates to optimize the user experience. This is challenging since the available network bandwidth (the exogenous input process) is hard to predict accurately. We simulate real-world video streaming using public cellular network data (Riiser et al., 2013) and video with seven bitrate levels and 500 chunks (DASH Industry Form, 2016). The reward is a weighted combination of video resolution, time paused for rebuffering, and the number of bitrate changes (Mao et al., 2017). The observed state contains bandwidth history, current video buffer size, and current bitrate. The action is the next video chunk’s bitrate. State-of-the-art heuristics for this
problem conservatively estimate the network bandwidth and use model predictive control to choose the optimal bitrate over the near-term horizon (Yin et al., 2015).

heading: K EXPERIMENT DETAILS
text: In our discrete-action environments (§6.2), we build 10-value networks and a meta-baseline using MAML (Finn et al., 2017), both on top of the OpenAI A2C implementation (Dhariwal et al., 2017). We use γ = 0.995 for both environments. The actor and the critic networks have 2 hidden layers, with 64 and 32 hidden neurons on each. The activation function is ReLU (Nair & Hinton, 2010) and the optimizer is Adam (Chilimbi et al., 2014). We train the policy with 16 (synchronous) parallel agents. The learning rate is 1−3. The entropy factor (Mnih et al., 2016) is decayed linearly from 1 to 0.001 over 10,000 training iterations. For the meta-baseline, the meta learning rate is 1−3 and the model specification has five step updates, each with learning rate 1−4. The model specification step in MAML is performed with vanilla stochastic gradient descent.
We introduce disturbance into our continuous-action robot control environments (§6.1). For the walker with wind (Figure 1c), we randomly sample a wind force in [−1, 1] initially and add a Gaussian noise sampled from N (0, 1) at each step. The wind is bounded between [−10, 10]. The episode terminates when the walker falls. For the half-cheetah with floating tiles, we extend the number of piers from 10 in the original environment (Clavera et al., 2018a) to 50, so that the agent remains on the pathway for longer. We initialize the tiles with damping sampled uniformly in [0, 10]. For the 7-DoF robot arm environments, we initialize the target to randomly appear within (−0.1,−0.2, 0.5), (0.4, 0.2,−0.5) in 3D. The position of the target is perturbed with a Gaussian noise sampled from N (0, 0.1) in each coordinate at each step. We bound the position of the target so that it is confined within the arm’s reach. The episode length of all these environments are capped at 1,000.
We build the multi-value networks and meta-baseline on top of the TRPO implementation by OpenAI (Dhariwal et al., 2017). We turned off the GAE enhancement by using λ = 1 for fair comparison. We found that it makes only a small performance difference (within ±5% using λ = {0.95, 0.96, 0.97, 0.98, 0.99, 1}) in our environments. We use γ = 0.99 for all three environments. The policy network has two hidden layers, with 128 and 64 hidden neurons on each. The activation function is ReLU (Nair & Hinton, 2010). The KL divergence constraint δ is 0.01. The learning rate for value functions is 1−3. The hyperparameter of training the meta-baseline is the same as the discrete-action case.
L INPUT-DEPENDENT BASELINES WITH PPO
Figure 8 shows the results of applying input-dependent baselines on PPO (Schulman et al., 2017) in MuJoCo (Todorov et al., 2012) environments. We make three key observations. First, compared to Figure 4 the best performances of PPO in these environments (blue curves) are better than that of TRPO. This is as expected, because the variance of the reward feedbacks in these environments is generally large and the reward clipping in PPO helps. Second, input-dependent baselines improve the policy performance for all environments. In particular, the meta-learning approach achieves the best performance, as it is not restricted to a fixed set of input sequences during training (§5). Third, the trend of learning curve is similar to that in TRPO (Figure 4), which shows our input-
dependent baseline approach is generally applicable to a range of policy gradient based methods (e.g., A2C (§6.2), TRPO (§6.1), and PPO).
M INPUT-DEPENDENT BASELINES WITH RARL
Our work is orthogonal and complementary to adversarial and robust reinforcement learning (e.g., RARL (Pinto et al., 2017)). These methods seek to improve policy robustness by co-training an adversary to generate a worst-case noise process, whereas our work improves policy optimization itself in the presence of inputs like noise. Note that if an adversary generates high-variance noise, similar to the inputs we consider in our experiments (§6), techniques such RARL alone are not adequate to train good controllers.
To empirically demonstrate this effect, we repeat the Walker2d with wind experiment described in §6.1. In this environment, we add a noise (of the same scale as the original random walk) on the wind and co-train an adversary to control the strength and direction of this noise. We follow the training procedure described in RARL (Pinto et al., 2017, §3.3).
Figure 9 depicts the results. With either the standard state-dependent baseline or our input-dependent baseline, RARL generally improves the robustness of the policy, as RARL achieves better testing rewards especially in the low reward region (i.e., compared the yellow curve to green curve, or red curve to blue curve in CDF of Figure 9). Moreover, input-dependent baseline significantly improves the policy optimization, which boosts the performance of both TRPO and RARL (i.e., comparing the blue curve to the green curve, and the red curve to the yellow curve). Therefore, in this environment, the input-dependent baseline helps improve the policy optimization methods and is complementary to adversarial RL methods such as RARL.
N INPUT-DEPENDENT BASELINES WITH META-POLICY ADAPTATION
There has been a line of work focusing on fast policy adaptation (Clavera et al., 2018a;b; Harrison et al., 2017). For example, Clavera et al. (2018b) propose a model-based meta-policy optimization approach (MB-MPO). It quickly learns the system dynamics using supervised learning and uses the learned model to perform virtual rollouts for meta-policy adaptation. Conceptually, our work differs because the goal is fundamentally different: our goal is to learn a single policy that performs well in the presence of a stochastic input process, while MB-MPO aims to quickly adapt a policy to new environments.
It is worth noting that the policy adaptation approaches are well-suited to handling model discrepancy between training and testing. However, in our setting, there exists no model discrepancy. In particular, the distribution of the input process is the same during training and testing. For example, in our load balancing environment (Figure 1a, §6.2), the exogenous workload process is sampled from the same distribution during training and testing.
Therefore our work is conceptually complementary to policy adaptation approaches. Since some of these methods require a policy optimization step (e.g., (Clavera et al., 2018b, §4.2)), our inputdependent baseline can help these methods by reducing variance during training. We perform an
experiment to investigate this. Specifically, we apply the meta-policy adaptation technique proposed by Clavera et al. (2018b) in our Walker2d environment with wind disturbance (Figure 1c, §6.1). For this environment, although the wind pattern is drawn from the same stochastic process (random walk), we aim to adapt the policy to each particular instantiation of the wind.
Operationally, to reduce complexity, we bypass the supervised learning step for the system dynamics and use the simulator to generate rollouts directly, since the interaction with the simulator is not costly for our purpose and the state transition in our environment is not deterministic. Following the meta-policy adaptation approach, the policy optimization algorithm is TRPO (Schulman et al., 2015a). The meta-policy adaptation algorithm is MAML (Finn et al., 2017). In particular, we performed ten gradient steps to specialize the meta-policy for each instantiation of the input process. For input-dependent baseline, we inherit our meta-baseline approach from §5. Similar to policy adaptation, we adapt our meta-baseline alongside with the policy adaptation in the ten gradient steps for each input instance.
The results of our experiment is shown in Figure 10. The learning curve (left figure) shows the policy performance for 100 unseen test input sequences at each training checkpoint. We measure the performance of MPO after ten steps of policy adaptation for each of the 100 input sequences. As expected, policy adaptation specializes to the particular instance of the input process and improves policy performance in the learning curve (e.g., MPO improves over TRPO, as shown by the green and yellow learning curve). However, policy adaptation does not solve the problem of variance caused by the input process, since the policy optimization step within policy adaptation suffers from large variance. Using an input-dependent baseline improves performance both for TRPO and MPO. Indeed, MPO trained with the input-dependent baseline (and adapted for each input sequence) outperforms the single TRPO policy, as shown by the red learning curve.
This effect is more evident in the policy adaptation curve (right figure). The policy adaptation curve shows the testing performance of the adapted policy at each adaptation step (the meta-policy is taken from the 5e7 training timestep). With an input-dependent baseline, the meta policy already performs quite well at the 0th step of policy adaptation (without any adaptation). This is perhaps unsurprising, since a single policy (e.g., the TRPO policy trained with input-dependent baseline) can achieve good performance in this environment. However, specializing the meta-policy for each particular input instance further improves performance.

Review summaries:
This work considers the problem of learning in input-driven environments -- which are characterized by an addition stochastic variable z that can affect the dynamics of the environment and the associated reward the agent might see. The authors show how the PG theorem still applied for a input-aware critic and then they show that the best baseline one can use in conjecture with this critic is a input-dependent one. My main concerns are highlighted in points (3) and (4) in the detailed comments below.