Truly Batch Apprenticeship Learning with Deep Successor Features

1 Introduction

Reward design is a key challenge in Reinforcement Learning (RL). Manually identifying an appropriate reward function is often difficult, and poorly specified rewards can pose safety risks [Leike et al., 2017]. Apprenticeship learning is the process of learning to act from expert demonstrations. One type of apprenticeship learning, Imitation Learning (IL) (e.g. [Ross et al., 2011; Ho and Ermon, 2016]), directly learns a policy from demonstrations. In contrast, Inverse Reinforcement Learning (IRL) aims to recover the expert policy by learning a reward function, which, when solved, induces policies that are similar to the expert. While there exist theoretical connections between IL and IRL [Piot et al., 2017], one may be preferred depending on the application. In particular, directly learning a policy (imitation) can be brittle in cases of long-horizon planning, and environments with strong covariate or dynamics shifts [Piot et al., 2013; Fu et al., 2017]. Besides addressing these issues, the learned reward function in IRL can also be used to identify expert motivations underlying the actions: rewards describe what

∗ Supplementary material (Appendix) to the main text can be found here : https://tinyurl.com/ijcai-dsfn-supplement

the expert wishes to achieve, rather than simply what they are reacting to, enabling agents to generalize better with the knowledge of these "intentions" in related environments.

In this work, we focus on IRL in batch settings: we must infer a reward function that induces the expert policy, given only a fixed set of expert demonstrations. No simulator exists and environment dynamics (transition model) is unknown. Performing analyses on batch data often ends up being the only reasonable alternative in domains such as healthcare, finance, education, or industrial engineering where precollected logs of expert behavior are relatively plentiful but new data acquisition or a policy roll-out is costly and risky.

Many existing IRL algorithms (e.g. [Abbeel and Ng, 2004; Ratliff et al., 2006]) use the feature expectations of a policy as the proxy quantity that measures the similarity between expert policy and an arbitrary policy that IRL proposes. If a simulator is available, feature expectations can be computed by taking sample means across on-policy rollouts [Abbeel and Ng, 2004]. However, new rollouts are not possible in batch settings. To estimate feature expectations in batch settings, a few IRL algorithms exist that either use a linear estimator [Klein et al., 2012] or bypass the estimation by assuming the existence of additional data from an arbitrary policy [Boularias et al., 2011]. Often, linear estimators do not possess the representational power necessary to model real-world tasks, while the ability to access additional data is overly restrictive. Along with an off-policy estimator for feature expectations, successful batch IRL also requires methods to engineer expressive feature spaces and manage computational complexity. Our work addresses all these batch IRL challenges.

Contributions Specifically, our work makes two key contributions that together enable a batch version of max-margin IRL that empirically scales well across classical control and real-world tasks, using only expert demonstrations and no additional inputs whatsoever. Firstly, we propose the Deep Successor Feature Network (DSFN), an off policy feature expectations estimator that replaces the on-policy roll-outs in standard IRL algorithms. Secondly, to mitigate a potential high bias caused by data support mismatch in DSFN, we propose Transition Regularized Imitation Learning (TRIL), which warm-starts the IRL loop with a near-expert initial policy and simultaneously provides an expressive feature space for better IRL performance directly learned from the batch data without any manual feature engineering.

2 RelatedWork

This paper focuses on batch IRL, where the IRL evaluation metric (for evaluating the IRL policy's closeness with respect to expert) is feature expectations [Abbeel and Ng, 2004]. Our goal is to learn a reward function for imitating and interpreting the expert policy under unknown dynamics with expert demonstrations alone. This batch definition differs from other settings such as RE-IRL [Boularias et al., 2011] that requires access to non-expert data from an arbitrary policy for importance sampling or DM-IRL [Burchfiel et al., 2016] that needs additional demonstration queries made to the expert. Finally, while we adopt the max-margin IRL framework [Abbeel and Ng, 2004] in this work, our key contributionsTRIL, DSFN- are generic and can be applied across a broad class of IRL methods (e.g. probabilistic IRL methods such as [Ziebart et al., 2008]).

Batch Feature-Expectations IRL. A major challenge in batch IRL is estimating feature expectations. [Klein et al., 2011] made an important observation that this problem is similar to off-policy policy evaluation and proposed an estimator based on LSTD-Q [Lagoudakis and Parr, 2003]; this estimator suffers from the high sensitivity and weak representational power of linear systems. To address the computational complexity aspect of IRL, [Klein et al., 2012] proposed SCIRL that takes a supervised learning (classification) approach where feature expectations are estimated by a simple sample mean. SCIRL assumes a deterministic expert and relies on the heuristic assumption that feature expectations of non-expert actions can be approximated by a multiplicative constant factor of the feature expectations for expert actions (requires tuning for every domain). In contrast, our method is scalable across domains with little tuning required and only requires the expert demonstrations without any additional assumptions. Moreover, SCIRL uses a linear classifier while our method admits any parametric model (e.g. neural nets). Our flexible neural net parametrization for estimating feature expectations in DSFN (similar in spirit to [Kulkarni et al., 2016]'s deep successor features for value functions in RL) along with a TRIL warm-start enables our model to scale the best among alternatives, to real-world tasks.

Warm-starting batch IRL. Computational complexity and sensitivity to features are unavoidable practical challenges in IRL, more so in batch settings in which offpolicy evaluations are required to compute feature expectations. Standard IRL algorithms [Abbeel and Ng, 2004; Ratliff et al., 2006] assume manually engineered features and initialize with a simple (random, sample mean) policy. However, in batch settings, initializing with such simple policies is not helpful as OPE on such policies can produce biased feature expectations estimates when expert actions may differ significantly from the policy itself. Our novel TRIL network warm-starts batch IRL to obtain a near-expert initial policy via supervised classification (similar ideas include [Klein et al., 2013; Piot et al., 2014]) regularized by the joint learning of dynamics, and simultaneously performs deep feature representation learning with the network's shared layers (e.g. [Song et al., 2016]).

3 Background

Markov Decision Process (MDP). An MDP is a tuple (S,A, T , T0,R, γ) where s ∈ S states (continuous, in this work), a ∈ A actions (discrete, in this work), T (s′ |s, a), T0 the transition probabilities and the initial state distribution respectively, R(s, a) the reward function, and γ ∈ [0, 1) the discount factor. A policy π(a|s) gives the probability of taking an action a in a state s. The state-value function is defined as V π (s) = Eπ[∑∞t=0 γ tR(st, at)|s0 = s]. The action-value function is defined as Qπ (s, a) = R(s, a) + Es1,a1∼π,...[∑∞t=1 γ tR(st, at)]. The optimal policy under an MDP is given by πe = π∗ = arg maxπ V π (s) (∀s ∈ S).

Batch Max-Margin IRL. We assume the existence of an expert policy πe that is optimal under some unknown, linear reward function of the form R(s, a) = w · φ for some reward weights w ∈ Rd and some pre-defined feature map φ(s, a) : S × A → Rd . Both quantities are bounded i.e. ||w||2 ≤ 1, ||φ(·)||2 ≤ 1. Let De = {(s0, a0 ∼ πe, . . . , sT )} be a set of N expert trajectories sampled according to πe. Unlike traditional IRL methods, we assume we cannot sample trajectories from any other policy and that T is unknown (common in real-world batch settings). The state-action feature expectations of a policy µπ (s, a) ∈ Rd and the overall feature expectations µπ ∈ Rd of a policy are defined as

Conceptually, the feature expectation µπ represents the (expected, discounted) amount of the feature accumulated while acting under a policy π. All IRL algorithms require a way to measure the similarity between the expert policy and candidate policies, and max-margin IRL uses ‖µπe - µπ ‖2 to measure the same. Specifically, [Abbeel and Ng, 2004] showed that convergence in feature expectations implies convergence in the expected value function between two policies.

4 Method: Batch Feature-Expectations IRL

In this section, we present our model-a batch version of max-margin IRL [Abbeel and Ng, 2004] that consists of two novel neural-network models: one designed to address the off-policy evaluation challenge intrinsic to batch settings and the other to warm-start the former and reduce both the biased estimations of OPE and IRL's computational complexity.

The main roadblock in batch IRL is the inability to estimate feature expectations by taking a sample mean of onpolicy rollouts: µπ ̸≈ 1N ∑Ni=1 ∑Tt=0 γ tφ(st, πe(st))) due to the lack of a simulator. If a candidate policy π disagrees with the expert demos such that at ∼ πe ̸= π(st) for some (st, at, st+1) ∈ De, we have no way of collecting next state samples to know how the trajectory evolves from st using the demonstrations alone. Computing the feature expectations now becomes an Off-Policy Evaluation (OPE) problem (e.g. [Thomas and Brunskill, 2016]) whose target is µπ instead of V π . In the following, we develop a model-free approach for the OPE, which avoids the bias of model-based methods and the high variance of importance sampling-based methods.

4.1 Estimating µπ via DSFN

Inspired by the LSTD-based approach [Klein et al., 2011], we observe an analogous formulation between feature expectations µπ ∈ Rd and action value function Qπ ∈ R. Note that µπ 's i-th component µπi (s, a) = φi(s, a) + Es1,a1∼π,...[∑∞t=1 γ tφi(st, at)] ∈ R, while Qπ (s, a) = R(s, a) + Es1,a1∼π,... [∑∞t=1 γ tR(st, at)] ∈ R. Thus, learning µπ can be setup as a system of d Temporal Difference (TD) learning problems [Sutton et al., 1998] which we can solve by assuming that each feature dimension µi is independently associated with its own reward function φi. Leveraging this similarity, we derive our method, Deep Successor Feature Networks (DSFN), in a way analogous to deep Q-learning algorithms such as DQN [Mnih et al., 2015]. Let µπ (s, a; θ) be the feature expectations estimator whose parameters θ (feedforward neural network) are to be learned from expert demonstrations De. The model is trained using gradient descent on MSE loss L(θ, π) with respect to the TD targets that follow from Bellman equation [Sutton et al., 1998]. Given π, we set the TD targets µπ+ ∀ (si, ai, s′i ) ∈ De, |De| = N as follows:

Algorithm (1) presents a batch version of max-margin IRL using DSFN. DSFN's training procedure (Appendix Section 2.1) is similar to DQN with a subtle difference: DSFN evaluates a policy while DQN optimizes and hence it has a max operator in its target. While typical max-margin IRL returns a set of policies and one must choose after convergence (no guarantee that the last solution is the best), we observed empirically in batch problems that using validation error of the Algorithm 1 Batch Max-Margin IRL

Input: π0 (Init. Policy), φ (Feature Map), De (Demo.), n (max. iteration) ϵ, δ (convergence thresholds)

Parameter: w, θ, Output: Rw

• 1: Estimate µπe by taking a sample mean over De
• 2: for i = 0 : n do
• 3: Estimate µπ(i) with µ(θ) DSFN and De
• 4: Obtain a reward function w(i) by solving QP:

• 5: Obtain π(i+1) by solving MDP with Rw = 〈w(i), φ〉.
• 6: if ∥∥∥µπe - µπ(i)∥∥∥2 < ϵ or Lval(π(i+1)) < δ then
• 7: Terminate the loop.
• 8: end if
• 9: end for
• 10: return Rw = 〈w(i), φ〉

IRL policy's actions compared to expert actions was a useful convergence criterion, particularly given that µπ computed via OPE could be biased. We take a validation set from De and terminate the algorithm when validation error Lval goes below some δ > 0. DSFN's parametric details are given in Appendix Table 2.

Necessity of warm-starting. As with any off-policy evaluation, our model-free DSFN approach to computing feature expectations can only be expected to be reliable if the candidate policy π has sufficient overlap or support under the transitions from the expert policy πe. Specifically, note that Eqn. (2) depends on (s, a, s′ ) ∈ De. If an evaluation policy π generates trajectories far from demonstrated trajectories where data are almost absent (i.e. D ∼ π, D∩De ≈ ∅), then DSFN's off-policy estimates can suffer from a high bias, for there would be no samples that characterize (s, a ∼ π(s)). One simple but important workaround is to ensure that π0 is already a near-expert policy, ensuring successive policy updates do not differ much from the expert. Starting near the expert policy also reduces the number of iterations of the IRL (each requires an MDP solution) loop till convergence.

4.2 Warm-starting DSFN-IRL with TRIL

The goal of our warm-start is to obtain good π0 and φ. It is known that a naive supervised learning approach tends to over-fit expert demonstrations due to correlated samples [Ross et al., 2011]. Our solution is to regularize the learned hidden layers by jointly requiring them to predict environmental dynamics in order to prevent over-fitting. Specifically, we propose Transition-Regularized Imitation Learning (TRIL), a two-channel network (Figure 2) that simultaneously obtains π0 via regularized multi-class classification and an empirically successful feature representation φ via deep representation learning in the process.

Regularized multi-class classification. Let θπ0 be the parameters for expert action prediction (upper channel in Fig. 2) and θT for next state prediction (lower channel in Fig. 2). The network is trained using the following loss: ∀(s, a, s′ ) ∈ De L(θπ0 , θT ) = LCE(a, π0(s; θπ0 )) + λLMSE(T (s, a; θT ), s′ ) where LCE denotes the cross entropy loss for expert action prediction given (s, a) ∈ De, LMSE the mean squared error

loss given (s, a, s′ ) ∈ De, and λ the regularization strength.

Deep representation learning. We observed empirically that the feature map φ learned from the shared layers (Figure 2 - notice that TRIL extracts only the hidden layers of θπ0 (shared across both channels) for obtaining φ(s)) of the TRIL network provides a significant performance boost to our batch IRL. Intuitively, the learned feature space extracts sufficient information from the current state to predict both expert actions and the next state dynamics in the environment, both of which we consider important to model the unknown rewards. Training a feature representation on these quantities avoids the need for expert-specified features, providing a scalable approach across tasks to learn good feature maps for IRL from the data itself. The feature learning aspect of TRIL model is inspired from the work of [Song et al., 2016] which states that feature encoders learned from dynamics prediction based errors are very successful value approximators in RL problems by demonstrating the theoretical connections between the model error for such a feature set and the Bellman error of the corresponding Q-function.

For discrete action problems (our experiments), φ(s, a) can be easily setup as the learned representation φ(s) from TRIL concatenated with a one-hot encoding that indicates the discrete action a i.e. φ(s, a) = concat(φ(s), {I(a = ai)}|A| i=1 ).

5 Experimental Setup

We evaluated our model (TRIL+DSFN) on two sets of tasks: three simulated classical control tasks and a real-world clinical management task. All experiments were done in a batch setting where we assumed that only the expert demonstrations were provided bereft of any ability to collect additional transitions or access the simulator while training. Models were tested for expert imitation and reward interpretability.

Evaluation and Model setup For the classical control tasks, we evaluated the learned policies in terms of cumulative rewards that the learned IRL policy acquires on the simulators after the independent batch IRL training. For sepsis management, we used action-matching validation accuracy to evaluate goodness of policies. Both TRIL and DSFN are multi-layer feed-forward neural networks with a Gaussian output layer. DDQN [van Hasselt et al., 2015] is used as the (near-optimal) MDP solver in our experiments. Parametric details of all the models are in Appendix Table 2.

Baselines. We compared our model to batch IRL baselines that use feature expectations explicitly and do not assume the existence of non-expert data: LSTD-µ [Klein et al., 2011] and SCIRL [Klein et al., 2012] (We did not include methods such as CSI, RCAL [Klein et al., 2013; Piot et al., 2014] as they do not use feature expectations; we leave comparison to a broader class of non-feature expectation-based IRL algorithms as future work). LSTD-µ is based on LSTD-Q [Lagoudakis and Parr, 2003] where it approximates the fixed points (µπ ) of the Bellman Equation by solving a linear system. SCIRL trains a linear multi-class classifier where expert feature expectations are computed using either a heuristic with Monte Carlo estimates or LSTD-µ. We chose LSTD-µ as we found it hard to apply the heuristic approach across tasks: the discount factor has to be tuned separately for each task and more importantly, it makes a relatively strong assumption that the effect of taking a non-expert action on its feature expectations can be expressed via the expert feature expectations up to the discount factor, which is limiting in practical batch setting tasks. Even with a tuned discount factor, a large, stochastic problem usually does not conform to the assumptions of the heuristic and the LSTD approach is more stable and generalizable across tasks with reasonably large amounts of data (e.g. sepsis). In general, we drew all experimental details from the authors' original source code (https://github.com/edouardklein/RL-and-IRL). To make the experiments fair, all batch IRL models had similar initializations, as appropriate. Link to our code repository can be found in Section 4 of the Appendix.

5.1 Classical Control Tasks

Task descriptions. MountainCar-v0, CartPole-v0 and Acrobot-v1 (https://github.com/openai/gym) are common simulated RL benchmarks. The environments are relatively small: MountainCar-v0 (S ⊂ R2 , |A| = 3), CartPole-v0 (S ⊂ R4 , |A| = 2), Acrobot-v1 (S ⊂ R6 , |A| = 3). We obtained near-optimal policies using DDQN on the simulators and generated demonstration data by varying the number of trajectories |De| = {1, 10, 100, 1000}. We note that while DDQN may not have produced an "optimal" expert, our goal in IRL is to simply identify rewards under which the observed behavior is optimal; our objective is not to recover the "true" reward that the expert may have been imperfectly targeting. We followed a 70-30 train-validation split in our batch data.

DSFN(+TRIL) outperformed baselines, had better data efficiency and performed comparably as the pure imitator (TRIL). Fig. (3) shows TRIL+DSFN (red) outperforms all batch IRL baselines, works admirably in low batch data regimes and scales well as the data sizes increase. We believe that LSTD-µ and SCIRL did not perform well in the low data regime because limited data led to an under-determined linear system and such LSTD models are known to be highly sensitive to data coverage (of the state-action space) of the expert demonstrations. Notably, the results also demonstrate that DSFN's action-matching performance is almost as good as pure imitator (TRIL) even when DSFN performs the harder and useful task of learning a reward function from demonstrations compared to pure policy imitation done by TRIL.

5.2 Clinical Task: Sepsis Management in ICU

Task descriptions. We evaluated TRIL+DSFN on a largescale real-world Sepsis management task with the goal of imitating and interpreting clinicians' demonstrations. Sepsis is a leading cause of cost and mortality in Intensive Care Units (ICU) [Mervyn et al., 2016]. Recently, [Raghu et al., 2017] performed deep RL to optimize a policy for intravenous fluids and vasopressor interventions for sepsis patients on a manually engineered reward function - an RL problem. Our work performs a complementary task: we take the same dataset and learn a reward function via batch IRL that helps imitate clinicians by inducing similar treatment policies. Learning such a reward function is valuable because it may expose components of the reward that experts forgot to code (e.g. ways to reduce morbidity as well as mortality; if patients cannot be saved, there may still be appropriate palliative actions to take). Second, we may expose aspects that experts are optimizing for unknowingly and unnecessarily, thus promoting behavior change. Thus, IRL provides a starting point for designing an appropriate reward function for future RL works.

Data and MDP formulation. The input data was obtained from the Multiparameter Intelligent Monitoring in Intensive Care (MIMIC-III v1.4) database [Johnson et al., 2016] and included a cohort of 17,898 patients fulfilling Sepsis-3 criteria. We then performed a 60-20-20 train-val-test split on our dataset. For the state space, we included 46 physiological features, including patient attributes, vitals and lab tests. For the action space, we included only the vasopressor intervention which was discretized to 5 bins (one for no action and four bins indicating 4 dosage quartiles) similar to [Raghu et al., 2017]. We included other interventions such as IV fluids, mechanical ventilations within the state space. Overall, this led to an MDP with S ⊂ R46 , |A| = 5 and unknown transition and reward functions. In-hospitality mortality indicator defined the terminal state ST . (More details on data, setup can be found in Appendix Section 3).

DSFN imitates clinicians better than the batch IRL baselines. Table (1) shows the action matching accuracy on the validation dataset where Top-{1, 3} measures the proportion of transitions in the test data in which the expert action is either the best or among the three best predictions of each method. We observe that our models (TRIL, TRIL+DSFN) outperform all baselines in the action-matching accuracy by ∼ 40%. Similar to control experiments, DSFN performs comparably as the pure imitator TRIL while undertaking the harder task of learning a reward function. For a fair comparison, we also warm-started the two batch IRL baselines with TRIL. In fact, when these baselines were not initialized with TRIL, their performance was even worse, likely due to the limited representation power of linear models and their sensitivity to batch data coverage (limited in clinical data). Notably, the poor performance of the unregularized supervised classifier (∼ 30% vs. ∼ 80% of TRIL) reinforces the value of transition-based regularization towards action imitation performance by preventing over-fitting to training batch data.

DSFN learns a clinically intuitive reward function. Fig. (4) analyzes the learned reward function with respect to three important vitals that are extreme in patients with Septic shock. Septic shock usually results in low platelet counts (PC), low blood pressures (BP), high heart rates (HR) and clinicians encourage strong dosages of vasopressors in these cases. We observed that the reward function complies with this intuition by giving low rewards for not administering vasopressor (dashed) and high rewards for administering vasopressor (solid) in the low BP, low platelets and high HR regimes. A practising intensivist confirmed this intuition.

6 Discussion

Performing IRL in a truly batch setting is challenging for three reasons: the estimation of feature expectations (specific to batch settings), IRL's computational complexity, and high sensitivity to the feature representation (generic across all IRL algorithms). In practice, it is essential to address all these challenges for batch IRL methods to perform consistently well across tasks. Our primary contribution, DSFN is key to address the OPE challenge for feature expectations and produces solid IRL performance compared to other baselines, even with the same initialization. TRIL, on the other hand, helps mitigate the other two challenges along with the potential data support mismatch issue of DSFN. With a TRIL warm-start, we typically converge in ≤ 5 iterations-significantly reducing computational burden-and the learned feature space from TRIL is expressive enough across tasks to find a reward function that induces policies almost as good as the expert.

In our efforts to solve these batch challenges, the role of few small, yet powerful, model engineering ideas cannot be overlooked. For instance, in cases where expert demonstrations are highly stochastic (such as the sepsis management), having an isotropic Gaussian output layer seems to provide a significant performance gain. The network learns the mean and variance of a Gaussian distribution that produces final samples similar to the work of [Duan et al., 2016]. We also normalize states on a rolling basis to provide a consistent range of input values to the networks [Henderson et al., 2017]. Similarly, transition-based regularization provides unprecedented performance boost on the imitation front by reducing over-fitting to batch data. In cases with highly correlated samples such as sepsis data, TRIL outperforms simple supervised imitation by about 50%. Besides, it is important to note that although TRIL learns a dynamics model, we use model training as a pure regularizer and our IRL algorithm remains model-free since the learned dynamics model (particularly for sepsis) is not powerful enough to simulate trajectories of IRL policies without accumulating huge errors.

As IRL is an ill-posed problem [Ng et al., 2000], the reward functions that our method (true for all IRL methods in general) learns is one among many reward functions that explain the observed demonstrations and no guarantee exists that the learned reward function matches the true (unknown) expert's reward function. Rather, the only possible assertion one can make is that the model learns a reward function that induces a policy whose value is sufficiently close to that of the expert. In future, one can look to restrict the possible class of reward functions to enhance reward identifiability [Fu et al., 2017].

Finally, we remark that TRIL and DSFN serve different purposes. If one only requires an expert-imitating policy under fixed dynamics, using TRIL alone may suffice. However, if one's focus is on understanding expert motivations via rewards or on building generalizable agents under shifting dynamics, learning rewards via DSFN (IRL) in batch settings is appropriate [Piot et al., 2013; Fu et al., 2017].

7 Conclusion

We proposed a novel IRL method that works well in batch settings without simulators and under unknown dynamics. Our method DSFN (+TRIL) outperforms all batch IRL baselines that use feature expectations on both simulated and realworld tasks, and the rewards learned in the ICU task match clinical intuition. Though we test our models on healthcare, our models make no domain specific assumptions and should work off-the-shelf for other domains. Also, the OPE and warm-start models we propose are generic and could be applied to develop batch versions of other popular IRL algorithms. Overall, our model is among the first ones to produce good imitation and learn rewards via IRL to infer the motivations underlying expert decisions in real-world batch settings.

Acknowledgements

We acknowledge and thank our funding source AFOSR FA 9550-17-1-0155.