Latent Hierarchical Imitation Learning for Stochastic Environments
Authors: Maximilian Igl
Punit Shah
Paul Mougin
Sirish Srinivasan
Tarun Gupta
Brandyn White
Kyriacos Shiarlis
Shimon Whiteson
Venue: Submitted to ICLR 2023 (2023)
Keywords: hierarchical imitation learning
learning from demonstrations
autonomous driving
causal confusion

Abstract: Many applications of imitation learning require the agent to avoid mode collapse and mirrorthe full distribution of observed behaviors. Existing methods improving this distributional realism typically rely on hierarchical policies conditioned on sampled types that model agent-internal features like persona, goal, or strategy. However, these methods are often inappropriate for stochastic environments, where internal and external factors of influence on the observed agent trajectories have to be disentangled, and only internal factors should be encoded in the agent type to be robust to changing environment conditions. We formalize this challenge as distribution shifts in the marginal and conditional distributions of agent types under environmental stochasticity, in addition to the familiar covariate shift in state visitations. We propose Robust Type Conditioning (RTC), which eliminates these shifts with adversarial training under randomly sampled types. Experiments on two domains, including the large-scal eWaymo Open Motion Dataset, show improved distributional realism while maintaining or improving task performance compared to state of the art baselines.

Under review conference paper at ICLR 2023
LATENT HIERARCHICAL IMITATION LEARNING FOR STOCHASTIC ENVIRONMENTS
Anonymous authors Paper under double-bling review
ABSTRACT
Many applications of imitation learning require the agent to avoid mode collapse and mirror the full distribution of observed behaviours. Existing methods that address this distributional realism typically rely on hierarchical policies conditioned sampled rypes that model are often inappropriate for stochastic environments, where internal and external factors of influence on the observed agent trajectories have to be disentangled, and only internal factors should be encoded in the agent type be robust to changing environment conditions. We formalize this challenge as distribution shift in the conditional distribution of agent types under environmental stochasticity, in addition to the familiar covariate shift in state visitations. We propose Robust Type Conditioning (RTC), which eliminate this shifts with adversarial training under randomly sampled types. Experiments on two domains, including the large-scale Waymo Open Motion Dataset, show improved distributional realism while maintaining or improving task performance compared to state-of-the-ar baselines.
features like persona, goal, or strategy. However, these methods
INTRODUCTION
Learning to imitate behaviour crucial when reward design is infeasible (Amodei et al., 2016: Hadfield- Menell et al. 2017 Fu et al. 2018 Everitt al. (2021), for overcoming hard exploration problems (Rajeswaran ct al. al. Zhu et al. 2018). y and for realistic modelling of dynamical systems with multiple interacting agents (Farmer and Foley 2009). Such systems, including games, driving simulations, and agent-based economic models, often have known state transition functions, but require accurate agents to be realistic. For example, for driving simulations, which are crucial for accelerating the development of autonomous vehicles (Suo al. 2021 [g] et al. 2022),, faithful reactions of all road users are paramount. Furthermore, is not enough to mimic a single mode in the data; instead, agents must reproduce the full distribution of behaviours to avoid sim2real gaps in modelled systems (Grover et al.) 2018 Liang et al.) 2020]. under-explored solutions in complex tasks (Vinyals et al., 2019) and suboptimal policies games requiring mixed strategies (Nash Jr 1950).
Current imitation learning (IL) methods fall short of achieving such distributional realism by matching all modes in the data. The required stochastic policy cannot be recovered from a fixed reward function and adversarial methods, while aiming to match the distribution in principle, are known to be prone to mode collapse in practice (Wang et al.] 2017]] Lucic et al., 2018: Creswell et al., 2018). Furthermore, progress on distributional realism is hindered by a lack of suitable IL benchmarks, with most relying on unimodal data and only evaluating task performance as measured by rewards, but not mode coverage. By contrast, many applications require distributional realism in addition to good task performance. For example, accurately evaluating the safety of autonomous vehicles in simulation relies on distributionally realistic agents. Consequently, goal is to improve distributional realism while maintaining strong task performance.
To mitigate mode collapse in complex environments, previous work uses hierarchical policies in an auto- encoder framework (Wang al.) 2017 Suo et al. 2021 et al. 2022). During training, an encoder infers

---

Under review as conference paper at ICLR 2023
latent variables from observed trajectories and the agent, conditioned on those latent variables, strives to imitate the original trajectory. At test time, a prior distribution proposes distributionally realistic latent values, without requiring access to privileged future information. We refer to this latent vector as an agent's inferred type since it expresses intrinsic characteristics of the agent that yield the multimodal behaviour. Depending on the environment, the type could, for example, represent the agent's persona, belief, goal, or strategy.
However, these hierarchical methods rely on either manually designed type representations (Igl et al. 2022) or the strong assumption that all stochasticity the environment can be controlled by the agent Wang ct al. 2017: Suo et al.) 202I). Unfortunately, this assumption is violated in most realistic scenarios. For example, in the case of driving simulations, trajectories depend not only on the agent's 's type, expressing its driving style and intent, but also on external factors such the behaviour of other road users. Crucially, despite being inferred from future trajectories during training, agent types must be independent of these external factors avoid leaking information about future events outside the agent's control, which in turn can impair generalization at test time under changed, and ex-ante unknown, environmental conditions. In other disentangle
words, the challenge in learning hierarchical policies using IL in stochastic environments the internal and external factors of influence on the trajectories and only encode the former into the type.
Consider the example of an expert approaching an intersection at the same time as another car. The expert passes if the other car brakes and yields to it otherwise. To reconstruct the scene with ease, a naively trained latent model could not only encode the agent's intended direction (an internal decision) but also whether to yield, which depends on the other car (an external factor). This is catastrophic at test time when the latent, and hence the yielding decision, is sampled independently of the other car's behaviour. In contrast, if only the expert's is intent were encoded in the latent, the policy would learn to react appropriately to external factors.
In this paper, we identify these subtle challenges arising under stochastic environments and formulate them new form of distribution shift for hierarchical policies. Unlike the familiar covariate shift in the state distribution (Ross et al. 201T), this conditional type shift occur in the distribution of the inferred latent type. It greatly reduces performance by yielding causally confused agents that rely on the latent type for information about external factors, instead of inferring them from the latest environment observation. We propose Robust Type Conditioning (RTC) to eliminate this distribution shift and avoid causally confused agents through a coupled adversarial training objective under randomly sampled types. We do not require access to an expert, counterfactuals, or manually specified type labels for trajectories.
Experimentally, we show the need for improved distributional realism due to mode collapse in state-of-the-art imitation learning techniques such as GAIL (Ho and Ermon 2016). Furthermore, we show that naively trained hierarchical models with inferred types improve distributional realism, but exhibit poor task performance in stochastic environments. By contrast, RTC can maintain good task performance in stochastic environments while improving distributional realism and mode coverage. We evaluate RTC on the illustrative Double Goal Problem as well as the large scale Waymo Open Motion Datasel (Ettinger et al. 2021 real driving behaviour.
2 BACKGROUND
drawn from p(T) of one or more experts interacting with a stochastic environment p(s++1 |8s; az) where 8, S are states and a, E A are actions. Our goal is to learn a policy T (a)a))))) to match p(T) when replacing the unknown expert and generat- ing rollouts f ~ p(f) p(50) - Ta ( 8,, âz) from the inital states 8o p(8)). We simplify notation and write ~ TT(()) and r ~ D(r) to indicate rollouts generated by the policy or drawn from the data respectively. Expectations E,>D and Ee~ are taken over all pairs (st, az) E and ($,, â,) f.
(6) We are given a dataset D = {Tp | of N trajectories T; = 80 ag ..... ST
Previous work (c.g., Ross et al. 2011 Ho and Ermon 2016) shows that a core challenge of learning from demonstration is reducing eliminating the covariate shift in the state-visitation frequencies p(s) caused by accumulating errors when using π.. Unfortunately, Behavioural Cloning (BC), a simple supervised training objective optimising maxg E,~D logg To (a, [s.]] is not robust to it. To overcome covariate shift, generative
2

---

Under review conference paper at ICLR 2023
adversarial imitation learning (GAIL) (o and Ermon 2016) optimises Te to fool a learned discriminator D (a,, 8,) that is trained to distinguish between trajectories in D and those generated by no:
min max EX-TT log( (a $.))] + B log(11 Da st))]
(1)
The policy can be optimised using reinforcement learning, by treating the log- dee-dscrriminttor scores as costs, rt -llog log ( $t). Alternatively, if the policy can be reparameterized (Kingma and Welling, 2013) and the environment is differentiable, the sum of log discriminator scores can be optimised directly without relying on high-variance score function estimators by backpropagating through the transition dynamics, Laav (f) Ei 12 log Dø(ār, &.)]. We refer to this Model-based GAIL (MGAIL), though in contrast to Baram et al. 2016), we assume a known differentiable environment instead of a learned model.
In this work, we are concerned with multimodal distributions p(T) and how mode collapse be avoided when learning Te. To this end, we assume the dataset is sampled from p(r) = p(sp) I ((g))(() Ir= |(ar (es.g))e ())((()+(()))))) ag d) where is the agent type, expressing agent char- acteristics such as persona, goal. or, strategy, and is a random variable capturing the stochasticity in the environment, i.e. p( (8,, Qt. E) is a delta distribution of gg) (8c+1) for some transition function f. Learned agents matching p(T), i.c., with p(f) ~ P(T), are distributionally realistic, whereas realism describes single trajectories when ê lies in the support of (r) As we show in section 6 current non-hierarchical adversarial methods (Ho and Ermon, 2016) exhibit mode collapse and are not distributionally realistic.
To combat mode collapse, hierarchical methods (e.g.. Wang et al., 2017: Lynch et al.) 2020 Suo et al. 2021 [gl et al.) 2022) often rely on an encoder to infer latent agent types 9. from trajectories during training, 9e ~ eo 9e IT. and optimise the control policy To To (ā, &4, 9 to generate trajectories f. similar T: fe ~ p((f.l are not accessible during testing, prior pe ( can be used to sample distributionally realistic types 9,. We indicate by subscript 9, or 9. whether the inferred type and trajectory are drawn from the prior distribution po(Op) or encoder ee (T) Subscripts are omitted for states and actions to simplify notation. Inferred types and predicted trajectories without subscripts indicate that either sampling distribution could be used. For discussing information theoretic quantities, we will use capital letters S, A, Å, G and 3 to denote the random variables for values s, a, à, 9 and E.
p( 80) II=i0 Te(ā,, (āg 82, 9e )(( |ū,, 82), with 9. (T) If ground truth trajectories
3 CONDITIONAL TYPE SHIFT IN STOCHASTIC ENVIRONMENTS
In this section we outline a challenge that arises for hierarchical policies in stochastic environments. A shift in the conditional type distribution can arise because latent types are drawn from the encoder éø(ğe (T) during training but from the prior pe (g,))) during testing. While the prior is trained to match the marginal distribution of the encoder p(..) = EE & (e [7)], this is not the case for the conditional distribution p(.. k)) 1g ( [T).. We show that this conditional type shift can result in policies ignoring environmen- tal information. In section o we experimentally confirm that this translates to reduced task performance.
We use the simplified model in fig. to describe the consequences of the conditional type shift. This model has two sources of randomness in the data D: the environmental noise & and the multimodal type g of the expert we are mimicking. The crucial difference between g and & is that & represents external factors that the agent cannot control but to which it has to react, while encodes agent-interndl decisions that can be taken independently of &. In this simplified model, the state 8 is a deterministic function of only E as we disregard cross-temporal dependencies. During training. when real trajectories 7 = (s, a) are available, the inferred type 9. is drawn from the encoder E (T) During testing, without access to T, a prior po (.,)) is used. Actions à are drawn from the learned control policy Te 8, 9) and optimisation is performed to minimise a reconstruction loss Lree (a, â). We express the core result of the policy 'ignoring environmental information' using mutual information ((S,A), entropy H (A, G.) and conditional entropy H (AAG..)
3

---

Under review conference paper at ICLR 2023
s(E)
8(E)
Pp(s,a]
Q1
((05), g~ B[1 (11-))
a
9e
à
9P
à
80
a()) t0 + -)81
(a, à)
81
6
1 a(s;,)) ga, (1 ggaj
(a) Encoder: 9. c0(1
(b) Prior: 0, Po
(c) Example dataset.
Figure 1: Simplified, non-temporal setup with environmental noise & and multi-modallty induced by the unohserved agent type g. We denote r (s, a). The inferred type I is sampled from (0(r.l7) during training (left) and po (.p)) otherwise (middle). The control policy is [ g). Circles are random variables and squares deterministic functions. The loss C(a, à) penalises differences between a and â. Right: Example data, B denotes Bernoulli distributions.
Theorem 1. We assume the model pa po(ā s, a) = [e la, 8) 8) @[lss 9.) ) : is achieving optimal re- (ā|s, (..l......d. 9,, with the
construction loss free on Po(s, a). The test policy po(ā s) = (g,),)) π (ā marginal encoder po po(ğ) (g) Ep, [ a. s)] as prior distribution. We can say for the training distribution P(s, a. 9.) = Pp(s, (). a) and and testing distribution P(s, â a. = P Pp(s)pol (g,, To s, 9,)): H(A < I(S, A) and (AA, G.) = (Á,, G,), then [(S, A) < I(S. A). Corollary 1. If (AA G.) = 0, the assumption H(A, G.) = H(, G,) becomes unnecessary in theorem and we have (5, Á) = 0 I(S, A).
Proofs are in appendix A The core result is that the mutual information between states and actions lower for prior policies than in the data, implying that such policies ignore action-relevant information in the states. This happens if H HAAGa.) < T(S, A), i.e. if the type Ò captures too much information about á during training. The condition H(A, G.) = H(Á, Gp)) assures that the entropy in the system and mutual information between variables remains comparable between training and testing. In the extreme case that type fully determines the action, i.e. HAA = 0, the policy ignores the state entirely, i.e. I(S, A) = 0.
Because the encoder and policy are trained jointly, the failure case H(A[G.) < T(S, A) requires the encoder capture too much information about in 9. and the policy relying too much on 9. to predict a, which constitutes a form of causal confusion. Without the encoder providing excessive information about the policy could not learn to over-rely on -.. Conversely, even with excessive information about 9.. the policy could still ignore and avoid H H(A[..) < I(S,A) A).
As an example, for the data given in fig. Ic this a solution satisfying A(S ((( 0:
0 =- if
implies Te(als.. 9) ag
if g 0 and pe(ğ,) B(005). (2)
(s,, a,)
a
if 9
with Bernoulli distribution B. The latent type fully determines the action and the policy ignores the state. This allows perfect reconstruction during training, but fails test time when po (T randomly sample au and a₁ with equal probability. In the example of a car approaching an intersection, this corresponds entering the intersection independently of whether another car is approaching quickly, leading to increased collision rates. Also note that the conditional type distribution is not independent of the state, i.e. plg. (o) B(1 <) + p(ğ. (s)) B(e), leading conditional type shift po(ğ,) B(0.5) test time.
would
Other training solutions exist for which I(S, A) H (Á(G) and I(S, A) = I(S, A). For example:
{1
10 ifi =
a;
if g
(s: a,)
To(als;.g)
for i + j. po (r B(1 c). (3)
if t
ag
if g
4

---

Under review conference paper at ICLR 2023
Ca,
Encoder IB
9c
Leve
ca(āelr)
Control policy Te(ā;) $t, 9)
pprion Colpo)
Prior Po(Op)
9p
Tp
Discriminator Dz(āe, &c)
Lade(f)
0...T
Figure 2: Robust Type Conditioning (RIC): The control policy we (a, (., Q) is trained under inferred types 9 sampled from both the encoder eo 29.911 IT) and the prior po p0()). The reconstruction loss L fe) avoids mode collapse. The adversarial loss Lav (fp) under prior types prevents causally confused policies and ensures good task performance. Cyoior optimises the prior to sample distributionally realistic types and the information bottleneck loss La reduces covariate shift.
Here the latent type 9 only captures the agent-internal randomness, the conditional type distribution matches the prior, i.e. p(..11) 81) (8) pe ( BB1 <), and the test policy correctly reproduces the data.
For temporally extended data, the states 8; will depend not only on E, but also on g or 9. complicating theoretical treatment. Nevertheless, seeing & as all future stochasticity in the environment, the same threat of conditional type shift arises. In the next section, we introduce two training interventions. The first discourages causally confused policies, the other discourages the encoder from capturing excessive information about &. section we show that using both interventions jointly allows to use hierarchical policies for improved distributional realism while avoiding the sub-optimal solution for which T(S, A) > T(S, A).
6
4 ROBUST TYPE CONDITIONING
We present Robust Type Conditioning (RTC), method for improving distributional realism in imitation learning while maintaining high task performance. RTC follows the auto-encoder framework discussed in sections 2 and but avoids conditional type shifts and causally confused policies that ignore environmental information in stochastic environments. This is achieved through two augmentations. First, during training, latent types are not only sampled from the encoder, but also the prior. Because we do not have ground truth trajectories for these prior sampled types, an adversarial loss is used in place of the reconstruction loss. Second, we regularise the mutual information T(S, G') using a variational information bottleneck (Alemi al. 2016) avoid excessive information in 9. RTC combines four losses: the reconstruction loss free, the information bottleneck loss Lib, the adversarial loss Ladv. and the prior loss Cprior (see fig. 2
LaTe (0 (0) ...) Cee (T, (T) + BL + Carr(Fe) peior (T)
(4)
(ED()))o(0))
Aadv (Tp) (
(7)].
Po(ū,) is a learned prior and (+(g) is shorthand for generating trajectories by rolling out the learned control policy π ( (. 9) in the environment. Parameters 0 are held fixed and A and / are scalar weights.
We now introduce the individual terms. First, Lee is reconstruction loss between and t.. encouraging the hierarchical policy to be distributionally realistic and encode useful information about the trajectory in the inferred type Ô. The loss Liee can take different forms. For example, in section 6. we use the BC loss (T) -llog log no (a,1 9.) while in section 6.2 we minimise the Le distance between agent positions in 8; and &, Note that state-based losses like the L2 reconstruction loss require access to a training environment able resimulate the conditions of the original trajectory т as assume that is still approximately optimal. The loss (Tri = Egg_ _e)( [log Po g..) optimises the prior to propose distributionally realistic types by matching the marginal encoder distribution.
5

---

Under review conference paper at ICLR 2023
The key algorithmic contribution of RTC is to also optimise the policy under types sampled from the prior (second line in eq. (.), not only the encoder (first line in eq. reconstruction loss cannot be used as the correct ground truth trajectories are unavailable. Instead, we use the advensarial loss Ladv (f) >. llo ( 82), where (a, 8t) is a learned discriminator (see section 2 This reduces conditional type shift because the prior distribution is already used during training. also reduces causal confusion: Because some types are now sampled independently of the trajectory and hence e. their information about& { is now less reliable and the policy is incentivised to rely on & as much as possible.
For these prior-sampled types, the
One can view sampling from the prior as causal intervention do(g) in which 9 is changed independently of the environmental factor ç. De Haan et al. (2019) show that causal confusion can be avoided by applying such interventions and optimising the policy correctly predict the counterfactual expert trajectory distribution, in our case Pespeat (Tls, do(g) ())). Unfortunately, we do not have access to this counterfactual trajectory. Instead, we rely on the generalisation of T to get us 'close' such a counterfactual trajectory for types dolğ) and then refine the policy locally using the adversarial objective.
We experimentally found that optimising the policy under prior types is sufficient to improve task performance and avoid causal confusion in hierarchical policies. However, it also eliminated improvements in distributional realism gained through the use of hierarchies, likely because the policy simply learned ignore the latent type altogether. As solution, we employ a informational bottleneck on types 9. ~ ( This filters information about & while still encoding information about g in 9.. thereby making the information in 9. more reliable and useful to the policy. We experimentally show that this, when combined with prior- type sampling during training, achieves improved distributional realism while maintaining excellent task performance. Without prior-type sampling, pask-performance degrades considerably, indicating that the bottleneck is insufficient for filtering out information about & entirely.
The information bottleneck preferentially filters information about &. because the control policy we also has direct access to it through the visited states 8. By contrast, information about g can only be accessed by the policy through 9. and is hence preferably encoded in the bottleneck when information bandwidth costs are applied. We found that both continuous type representations with Lb = KL[p(å (p(.)) ([: I)] and discrete type representations using straight straight-through gradient estimation work well in practice (see section 6.2..
To accommodate optimisation under inferred types drawn from both the encoder eg and the prior po. we split each minibatch B = (r(0) of N₂ trajectories sampled from 'D into two parts. For the fraction of of trajectories in B the rollouts è. are generated from types sampled from the encoder 9. ~ c0(rā.l) and all
four losses are optimised (first line in eq. (4). For the remaining fraction (1 - ʃ) of trajectories types are sampled from the prior pe(ý,) and only Lad and Lprior are optimised (second line in eq. (4)).
Optimisation of Laty and fre: can either be performed directly, similar to MGAIL. (Baram et al. 2016). by using a differentiable environment and reparameterised policies and encoder (Kingma and Welling 2013 or by treating them as rewards and using RL methods such as TRPO (Shhulman et al. 2015 Ho and Ermon 2016) or PPO (Schulman et al. 2017). The losses Ppiior and Lub can always be optimised directly.
5 RELATED WORK
Several previous works combine adversarial training with autoencoder architectures in the image domain. Makhzani et al. (2016) use adversarial loss on the latent variable in place of the KL-regularization used in VAËS. However, this eliminates the information bandwidth regularization for continuous latents which we show to be important for hierarchical imitation learning. Larsen et al. (2016) aim to learn similarity metric for visual inputs using latent representations of the discriminator. This is valuable for imitation learning from raw images (Rahmatizadeh al. 2018). but is not required for our experimental domains. Lastly, Chrysos jet (2018), similar to our work, use an additional autoencoding loss to better capture the data distribution in the latent space. However, they consider denoising images instead of imitation learning under stochasticity.
6

---

Under review conference paper at ICLR 2023
Hierarchical policies have been extensively studied in RL (e.g.. Sutton et al.) 1999: Bacon et al. 2017 et al.
2019 Igl et 2020) and II. In RL, they improve exploration, sample efficiency and fast adaptation. By contrast, in IL, hierarchies are used to capture multimodal distributions, improve data efficiency (Krishnan et al. 2017 1 Le et al. 2018). and enable goal conditioning (Shiarlis et al. 2018). Similar to our work, Wang et al. (2117) and Lynch et al. (2020) learn to encode trajectories into latent types that influence a control policy. Crucially, both only consider deterministic environments and hence avoid the distribution shifts and unwanted information Ieakage we address. They extend prior work in which the type. context, is provided in the dataset (Merel al. 2017). which is also assumed in (Fei et al. 2020). Tamar jet al. (2018) use a sampling method to infer latent types. Khandelwal et al. (2020) and Igl et al. (2022) use manually designed encoders specific to road users by expressing future goals as sequences of lane segments.
Vezhnevets et al.,
This avoids information leakage but cannot express all characteristics of human drivers, such as persona, and cannot transfer to other tasks. Futures states in deterministic environments (Ding et al. 2019), language PPashevich et al. 202T), and predefined strategy statistics (Vinyals et al. (019)) have also been used as types.
Information theoretic regularization offers an alternative learning hierarchical policies using the auto- encoder framework (Li et al. 2017: Hausman et al. (2077). However, these methods are less expressive since their prior distribution cannot be learned and only aim to cluster modes already captured by the agent but not pe- nalize dropping modes in the data. This provides useful inductive bias but often struggles in complex environ- ments with high diversity, requiring manual feature engineering (Eysenbach et al.) 2019 |Pathak et al., 2019).
Lastly, TrafficSim (Suo et al., 2021 uses IL to model driving agents and controls all stochasticity in the scene but uses independent prior distributions for separate agents. Hence, while no conditional distribution shift in p(ğlk) can occur (as E is constant), distribution shifts in p(g (( ()). and hence the joint marginal p( (1: can occur for latent types , of agents i * j with i.j E [1....} and ((I:N) = [gx 9
(N);
when drawn from the encoder, goals gf¹) and g0و) are coordinated through conditioning on the joint agent future, while they are independent when drawn from the prior. They use a biased "common sense" collision avoidance loss, motivated by covariate shift in visited states. Our work suggests that marginal type shift might also explain the benefits gained. In contrast, our adversarial objective is unbiased. See appendix [ for more related work on agent modelling in multi-agent settings, behavioural prediction and causal confusion.
6 EXPERIMENTS
We show in two stochastic environments with multimodal expert behaviour that i) existing adversarial methods suffer from insufficient distributional realism, ii) existing hierarchical methods cannot achieve good task performance and distributional realism and iii) RTC improves distributional realism while maintaining excellent task performance. We discuss differences in realism, coverage and distributional realism in fig, S
We compare the following models: MGAIL uses a learned discriminator and backpropagates gradients through the differentiable environment. It also optimises a BC loss as we found this to improve performance.
Symphony ([gl et al. 2022), building on MGAIL, utilises future lane segments as manually specified types (see appendix D.3). Our implementation of Symphony outperforms the results from (Igl et al. 2022) due to the additional use of a value function. InfoMGAIL Li et al. 2017) augments MGAIL to elicit distinct trajectories for different types. This introduces inductive bias but does not directly penalise mode collapse. Our methods, RTC-C and RTC-D, use a continuous or discrete type respectively. We also perform the ablation Hierarchy-NoPT (No Prior Training) which only uses the first line in eq. (4). i.e. f = 1. Hierarchy-Noor is similar to existing hierarchical methods, such as the proprietory TrafficSim (Suo) et al. 2021), in that it learns the prior but does not use it during training, only inference. It thereby does not account for distribution shifts in the latent types, as discussed in section 3
6.1 DOUBLE GOAL PROBLEM
In the double goal problem, the expert starts from the origin and creates a multimodal trajectory distribution by randomly choosing and approaching one of two possible, slowly moving goals located on the 2D plane.
7

---

Under review conference paper at ICLR 2023
Esper
KGNHL
RTC (o.rs)
Test Return
0.
JED Goal Distriburion
Frequency Lower Goal
Test Return (IB Abletion)
0.8
0.3
0.5
Bottlonack BCs
0.6
00
p=1.0
MGAIL
0.75
2-0.5
InfoMMAAA
0.2
0.50
9-0.6
BC
9-0.3
H-NoPT
0.0
0.3 0.00
8000
1300D
2000
10300
2500
4000
6000
UO00
Figure 3: Top: Visualization of ten randomly sampled goal pairs and associated trajectories. Bottom: Training curves, exponentially smoothed and averaged over 20 seeds. Shading shows the standard deviation. We show task performance as "Test Return' and distributional realism as 'ISD' between the goal distribution expert and agent (lower is hetter). "Frequency Lower Goal' shows the data from which the JSD is computed. The inset shows the distribution at the last training step. Boxes show quartiles, whiskers extreme values, diamonds outliers, and stars the mean.
Stochasticity introduced through randomized initial goal locations and movement directions. Nevertheless, the lower and upper goal (g) I remain identifiable by their location as yı < 0 for 9: and y/. > 0 for 9T, (see 3 While both goals are equally easy to reach, the expert has a preference P(G = 9:) 0.75. Sufficiently complex expert trajectories prevent BC from achieving optimal performance, requiring more advanced approaches. The expert follows a curved path and randomly resamples the selected goal for the first ten steps to avoid simple decision boundary along the a-axis in which experts in the lower half-plane always target goal 91- RTC uses the BC loss as reconstruction loss Lree (T) = log π (a. st 9.) and continuous types. All policies use bimodal Gaussian mixture model as action distribution. Performance measured the number of steps for which the agent is within 8 = 0.1 distance of one of the goals. We take h1 sign(yr) of the final agent position [IT, [rT to indicate the approached goal and measure distributional realism as the divergence between the empirical distributions. JSD (Pagem (
[PPexeet(h.)) (h,)). Details can be found in appendix D.1
Figure shows that MGAIL improves task performance compared to BC. Our method, RTC, improves it further, possibly because given a type, the required action distribution is unimodal. Importantly, RTC substantially improves distributional realism, achieving lower JSD values. To analyse this result, we show Pagent (h. = 1), 1). the frequency of targeting the lower goal. Not only is RTC's average value of PRre(h, --1) closer the true value of 0.75, it is also more stable across seeds, resulting in a lower JSD. The bias introduced by InfoMGA reduces task performance without improving distributional realism. As expected, the ablation Hierarch serrarchy-Noor NoPT achieves excellent distributional realism through the learned hierarchy but suffers reduced task performance due to unaccounted distribution shifts. Lastly, the rightmost plot of fig. shows that the information bottleneck is necessary.
6.2
WAYMO OPEN MOTION DATASET (WOMD)
To evaluate RIC on a complex environment we use the Waymo Open Motion Dataset (Ettinger et al. 2021 consisting of 487K segments of real world driving behaviour, Distributionally realistic agents are critical for driving simulations, for example for estimating safety metrics. Diverse intents and driving styles cause the data to be highly multimodal. Stochasticity is induced through the unpredictable behaviour of other cars, cyclists and pedestrians. We use (T, +) = E. (Steer(St. $t) where Liuber is the average Huber loss of the four vehicle bounding box corners. More details can be found in appendix D.2
8

---

Under review conference paper at ICLR 2023
Table 1: Averages and standard deviation over 10 training runs WOMD.
Collision
Off-road
MinADE
Curvature JSD
Progress JSD
rate (%) :
time (%) 4
(m)
(x10-9)
(x10-9) 9)
Data Distribution
1.16
0.68
MGAIL
5.39 0.68
0.89 0.12
1.34 + 0.08
1.32 - 1.48
3.81 1.29
Symphony
6.39 0.95
0.90 0.06
1.40 = 0.12
0.97 + 0.62
6,44 - 5.25
InfoMGAIL- C
5.21 + 0.37
0.89 0.14
0.07
4,40 1.47
InfoMGAIL D
4.82 0.29
0.84 0.10
1.15
0.77
4.01 1.45
Hierarchy-NoPT
35.08 + 0.44
1.83 0.42
1.12 + 0.01
1.76
2.54 0.63
RTC NolB
3,88 0.35
0.67 0.04
1.45 0.14
1.09 1.09::.06
2.84 0.65
RTC-
4.23 0.16
0.68 0.04
1.15 0.10
0.43 + 0.06
2.17 0.65
RTC D
4.21 + 0.24
0.74 + 0.06
1.12 0.10
0.89 + 0.66
2.56 = 0.54
We use the percentage of segments with collisions and time spent off-road as proxy metries for task performance. Mode coverage is measured by the minimum average displacement error, minADE (t,, ~wo min; Ei= 8(s ((s:,8,,) 85,0) where and we find the minimum over K 16 rollouts (hierarchical methods use K independently sampled types). Lower minADE DE implies better mode coverage, but does not directly measure the relative frequency of modes, e.g., low probability modes may be overrepresented. To measure distribution matching in driving intent, we use the Curvature JSD (Igl et al. 2022): in lane branching regions, such as intersections, it maps trajectories the nearest lane and extracts its curvature as feature heur. To compute JSD (Pagem (h [ |Pexpen (hear)), the value of hear is discretize into 100 equisized bins. To measure the driving style distribution, we extract the progress feature hstyle 8Sp, 8T) and use the same discretization to compute the JSD)
the Euclidean distance between agent positions
Results are provided in table Both versions of RTC improve task performance (collisions and off-road events) and distributional realism metrics (minADE and divergences) compared to the flat MGAIL baseline and previous hierarchical approaches (Symphony, Inf oMMAII, Hierarchy chyyN Both type represen- tations, RTC-C and RTC-D, perform similarly, showing robustness RTC to different implementations. The advantage of RTC in achieving both good task performance and distributional realism becomes clearest by comparing it to Hierar - and RIC-NOIB. While Hierarchy-ae -oPT achieves some improve- ments in distributional realism, is has nearly order of magnitude more collisions. This is consequence of the challenges discussed in section 3 which RTC is able to avoid. On the other hand, RTC-No IB, also avoids these challenges and achieves excellent task performance by using prior-sampled types during training. However, as discussed in section 4 it does not improve on distributional realism compared to flat baselines, indicating that the learned policy simply ignores the latent type. Combining prior-type type sampling and the information bottleneck achieves better distributional realism and task performance than all baselines.
7 CONCLUSIONS, LIMITATIONS, AND FUTURE WORK
This paper identified new challenges in learning hierarchical policies from demonstration to capture multi- modal trajectory distributions in stochastic environments. We expressed them conditional type shifts and causul confusion in the hierarchical policy. We proposed Robust Type Conditioning (RTC) to eliminate these distribution shifts and showed improved distributional realism while maintaining or improving task perfor- mance two stochastic environments, including the Waymo Open Motion Dataset (Ettinger et al., 2021). Future work will address conditional distributional realism by not only matching the marginal distribution p(r), but the conditional distribution p(r(E) (g) under specific realization of the environment. For example, drivers might change their intent based on the current traffic situation or players might adapt their strategy as the game unfolds. Achieving such conditional distributional realism will also require models and metries.

---

Under review conference paper at ICLR 2023
REFERENCES
A. A. Alemi, 1. Fischer, J. V. Dillon, and K. Murphy. Deep variational information bottleneck. arXiv preprint arXiv:1612.00410. 2016.
D. Amodei, C. Olah, J. Steinhardt, P. Christiano, J. Schulman, and D. Mané. Concrete problems in ai safety. arXiv preprint arXiv:1606.0656. 2016.
P.-L. Bacon, J. Harb, and D. Precup. The option-critic architecture. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 31, 2017.
N. Baram, O. Anschel, and S. Mannor. Model-based adversarial imitation learning. arXiv preprint 1611:1612 2016.
S. Bengio, O. Vinyals, N. Jaitly, and N. Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. Advances in neural information processing systems, 28, 2015.
S. Casas, C. Gulino, S. Suo, K. Luo, R. Liao, and R. Urtasun. Implicit latent variable model for scene- consistent motion forecasting. In European Conference on Computer Vision, pages 624-641. Springer, 2020.
Chai, B. Sapp, M. Bansal, and D. Anguelov. Multipath: Multiple probabilistic anchor trajectory hypotheses for behavior prediction. arXiv preprint arXiv:1910.05449, 2019.
G. G. Chrysos, J. Kossaifi, and S. Zafeiriou. Robust conditional generative adversarial networks. arXiv preprint arXiv: 08657, 2018
A. Creswell, T. White, V. Dumoulin, K. Arulkumaran, B. Sengupta, and A. A. Bharath. Generative adversarial networks: An overview. IEEE signal processing magazine, 35(1):53-65, 2018.
H. Cui, V. Radosavljevic, F-C. Chou, T.-H. Lin, T. Nguyen, T.-K. Huang, J. Schneider, and N. Djuric. Multimodal trajectory predictions for autonomous driving using deep convolutional networks. In 2019 International Conference on Robotics and Automation (ICRA), pages 2090-2096. IEEE, 2019.
De Haan, D. Jayaraman, and Levine. Causal confusion in imitation learning. Advances in Neural Information Processing Systems, 32, 2019.
Ding, C. Florensa, Abbeel, and M. Phielipp. Goal-conditioned imitation learning. Advances in neural information processing systems, 32, 2019.
Ettinger, S. Cheng, Caine, C. Liu, H. Zhao, Pradhan, Y. Chai, Sapp, C. Qi, Zhou, Z. Yang, A. Chouard, P. Sun, J. Ngiam, V. Vasudevan, A. McCauley, J. Shlens, and D. Anguelov. Large scale interac- tive motion forecasting for autonomous driving: The waymo open motion dataset. CoRR, abs/2104.10133, 2021. URL
htt.aa////211..... 10133
T. Everitt, M. Hutter, R. Kumar, and V. Krakovna. Reward tampering problems and solutions in reinforcement learning: A causal influence diagram perspective. Synthese, 198((7):6435-6)) 6467,, 2021.
B. Eysenbach, A. Gupta, Ibarz, and S. Levine. Diversity is all you need: Learning skills without reward function. In International Conference on Learning Representations, 2019. URL https 7/openreview.o .nt forum? 21d=211d33d2x3Fm
J. D. Farmer and D. Foley. The economy needs agent-based modelling. Nature, 460(7256):685-6688 -66)666,66666666 2009.
10

---

Under review conference paper at ICLR 2023
C. Fei, B. Wang, Y. Zhuang. Z. Zhang, J. Hao, H. Zhang, X. Ji, and W. Liu. Triple-gail: multi-modal imitation learning framework with generative adversarial nets. arXiv preprint arXiv:2005. 2020.
J. Fu, K. Luo, and S. Levine. Learning robust rewards with adverserial inverse inverse reinforcement learning. In International Conference on Learning Representations, 2018. URL httppss//oppeeeeiieeee
rkHywl-
A. Grover, M. Al-Shedivat, J. Gupta, Y. Burda, and H. Edwards. Learning policy representations in multiagent systems. In International conference on machine learning, pages 1802-1811. PMLR, 2018.
D. Hadeield-Menell, S. Milli, P. Abbeel, J. Russell, and A. Dragan. Inverse reward design. Advances in neural information processing systems, 30, 2017.
K. Hausman, Y. Chebotar, Schaal, G. Sukhatme, and Lim. Multi-modal imitation learning from unstructured demonstrations using generative adversarial nets. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, Fergus, Vishwanathan, and R. Garnett, editors, Advances in Neural Information Process- ing Systems, volume 30. Curran Associates, Inc., 2017. URL /ttp://t/////ccceiigg...... neurip cc st/ppaper/211/ff1ee 16663ccee3 6db83e7 22
Paper.
H. He, J. Boyd-Graber, K. Kwok, and H. Daumé III. Opponent modeling in deep reinforcement learning. In International conference machine learning, pages 1804-1813, PMLR, 2016.
P. Hernandez-Leal, B. Kartal, and M. E. Taylor. Agent modeling auxiliary task for deep reinforcement learn- ing. In Proceedings of the AAAI conference on artificial intelligence and interactive digital entertainment, volume 15, pages 31-37, 2019.
J. Ho and S. Ermon. Generative adversarial imitation learning. Advances in neural information processing systems, 29, 2016.
A. Hu, Z. Murez. N. Mohan, S. Dudas, J. Hawke, V. Badrinarayanan, R. Cipolla, and A. Kendall. Fiery: Future instance prediction in bird's-eye s-ed view from surround monocular cameras. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pages 15273-15282 2021.
M. Igl, A. Gambardella, J. He, N. Nardelli, N. Siddharth, W. Böhmer, and S. Whiteson. Multitask soft option learning. In Conference on Uncertainty Artificial Intelligence, pages 969-978. PMLR, 2020.
M. Igl, D. Kim, A. Kuefler, P. Mougin, P. Shah, K. Shiarlis, D. Anguelov, M. Palatucci, B. White, and S. Whiteson. Symphony: Learning realistic and diverse agents for autonomous driving simulation. arXiv preprint ar X3Xv:22055.03195, 05 2022.
B. Ivanovic and M. Pavone. The trajectron: Probabilistic multi-agent trajectory modeling with dynamic spatiotemporal graphs. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 2375-2384, 2019.
S. Khandelwal, W. Qi J. Singh, A. Hartnett, and D. Ramanan. What-if motion prediction for autonomous driving. arXiv preprint arXiv:2008. 10587, 2020.
D. Kingma and M. Welling. Auto-encoding variational bayes. arXiv preprint arXivv 2013.
Krishnan, R. Fox, Stoica, and K. Goldberg. Ddco: Discovery of deep continuous options for robot learning from demonstrations. In Conference on robot learning, pages 418-437. PMLR, 2017.
A. M. Lamb, A. G. ALIAS PARTH GOYAL, Y. Zhann, S. Zhang, A. C. Courville, and Y. Bengio. Professor forcing: A new algorithm for training recurrent networks. Advances in neural information processing systems, 29, 2016.
11

---

Under review conference paper at ICLR 2023
A. B. L. Larsen, S. K. Sønderby, H. Larochelle, and O. Winther. Autoencoding beyond pixels using a learned similarity metric. In International conference on machine learning, pages 1 1558-1566. PMLR, 2016.
H. Le, N. Jiang, A. Agarwal, M. Dudik, Y. Yue, and H. Daumé III. Hierarchical imitation and reinforcement learning. In International conference on machine learning, pages 2917-2926. PMLR, 2018
Li, J. Song, and S. Ermon. Infogail: Interpretable imitation learning from visual demonstrations. Advances in Neural Information Processing Systems, 30, 2017.
Y. Liang, C. Guo, Z. Ding, and H. Hua. Agent-based modeling in electricity market using deep deterministic policy gradient algorithm. IEEE Transactions on Power Systems, 35(6):4 (000:4180-192 2020.
Y. Liu, J. Zhang, L. Fang, Q. Jiang, and B. Zhou. Multimodal motion prediction with stacked transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 7577-7586, 2021.
M. Lucic, K. Kurach, M. Michalski, S. Gelly, and O. Bousquet. Are gans created equal? large-scale study. Advances in neural information processing systems, 31, 2018.
C. Lynch, M. Khansari, T. Xiao, V. Kumar, J. Tompson, S. Levine, and P. Sermanct. Learning latent plans from play. In Conference on robot learning, pages 1113-1132. PMLR, 2020.
A. Makhzani, Shlens, N. Jaitly, and I. Goodfellow. Adversarial autoencoders. In International Conference on Learning Representations, 2016. URL |RL http://axxivvoooooo abs /et.//11l... 05644
J. Merel, Y. Tassa, D. TB, S. Srinivasan, J. Lemmon, Z. Wang, G. Wayne, and N. Heess. Learning human behaviors from motion capture by adversarial imitation. arXiv preprint arXiv: 1707.0220, 2017.
O. Nachum, S. Gu, H. Lee, and S. Levine. Near-optimal representation learning for hierarchical reinforcement learning. In International Conference on Learning Representations, 2019. URL/cttps://pperreviie net/ Sttt/foomm?IG==Heemssooo
J. F. Nash Jr. Equilibrium points in n-person games. Proceedings of the national academy of sciences, 36(1): 48 49, 1950.
P. A. Ortega, M. Kunesch, G. Delétang, T. Genewein, J. Grau-Moya, J. Veness, J. Buchli, J. Degrave, B. Piot, J. Perolat, et al. Shaking the foundations: delusions in sequence models for interaction and control. arXiv preprint arXiv:21 10819, 2021.
G. Papoudakis and S. V. Albrecht. Variational autoencoders for opponent modeling in multi-agent systems. arXiv preprint arXiv:2001. 10829, 2020.
J. Park, Y. Seo, C. Liu, L. Zhao, T. Qin, J. Shin, and T.-Y. Liu. Object-aware regularization for addressing causal confusion in imitation learning. Advances in Neural Information Processing Systems, 14:33029-3044 2021.
A. Pashevich, C. Schmid, and C. Sun. Episodic transformer for vision-anddlanguaag navigation. In Proceed- ings of the IEEE/CVF International Conference on Computer Vision, pages 15942-15952, 2021.
D. Pathak, D. Gandhi, and A. Gupta. Self-supervised exploration via disagreement. In International conference machine learning, pages 5062-5071. PMLR, 2019.
T. Phan-Minh, E. C. Grigore, F. A. Boulton, O. Beijbom, and E. M. Wolff. Covemet: Multimodal behavior prediction using trajectory sets. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pages 14074-14083 2020.
12

---

Under review conference paper at ICLR 2023
R. Rahmatizadeh, P. Abolghasemi, L. Bölöni, and S. Levine. Vision-based multi-task manipulation for inexpensive robots using end-to-end learning from demonstration. In 2018 IEEE international conference on robotics and automation (ICRA), pages 3758-3765. IEEE, 2018.
R. Raileanu, E. Denton, A. Szlam, and R. Fergus. Modeling others using oneself in multi-agent reinforcement learning. In International conference on machine learning, pages 4257 4267-4266. PMLR, 2018.
A. Rajeswaran, V. Kumar, A. Gupta, G. Vezzani, J. Schulman, E. Todorov, and S. Levine. Learning complex dexterous manipulation with deep reinforcement learning and demonstrations. arXiv preprint arXiv:1709. 10087, 2017.
Ross, G. Gordon, and D. Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pages 627-635, JMLR Workshop and Conference Proceedings, 2011.
T. Salzmann, B. Ivanovic, P. Chakravarty, and M. Pavone. Trajectron++: Dynamically-feasibee trajectory forecasting with heterogeneous data. In European Conference on Computer Vision, pages 683-700. Springer, 2020.
J. Schulman, S. Levine, Abbeel, M. Jordan, and Moritz. region policy optimization. In International conference on machine learning, pages 1889-1897. PMLR, 2015.
J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal policy optimization algorithms. arXiv preprini arXiv: 1707:0007444 2017.
K. Shiarlis, M. Wulfmeier, S. Salter, S. Whiteson, and I. Posner. Taco: Learning task decomposition via temporal alignment for control. In International Conference on Machine Learning, pages 4654-4663. PMLR. 2018.
Suo, Regalado, S. Casas, and Urtasun. Trafficsim: Learning to simulate realistic multi agent behaviors. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10400-10409, 2021.
R. S. Sutton, D. Precup, and S. Singh. Between mdps and comi-mdps: A framework for temporal abstraction in reinforcement learning. Artificial intelligence, 1 (12(1-2):181 1111 1999.
A. Tamar, K. Rohanimanesh, Y. Chow, C. Vigorito, B. Goodrich, M. Kahane, and D. Pridmore. Imitation learning from visual data with multiple intentions. In International Conference on Learning Representations, 2018.
C. Tang and R. R. Salakhutdinov. Multiple futures prediction. Advances in Neural Information Processing Systems, 32, 2019.
A. S. Vezhnevets, S. Osindero, T. Schaul, N. Heess, M. Jaderberg, D. Silver, and K. Kavukcuoglu. Feudal networks for hierarchical reinforcement learning. In International Conference on Machine Learning, pages 3540-3549. PMLR, 2017
O. Vinyals, I. Babuschkin, W. M. Czarnecki, M. Mathieu, A. Dudzik, J. Chung, D. H. Choi, R. Powell, Ewalds, P. Georgiev, et al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 575(7782):350- 354, 2019.
Z. Wang, J. S. Merel, S. E. Reed, N. de Freitas, G. Wayne, and N. Heess. Robust imitation of diverse behaviors. Advances in Neural Information Processing Systems, 30, 2017.
13

---

Under review conference paper at ICLR 2023
C. Wen, Lin, T. Darrell, Jayaraman, and Y. Gao. Fighting copycat agents in behavioral cloning from observation histories. Advances in Neural Information Processing Systems, 33:2564 2575, 2020.
A. Xie, D. P. Losey, R. Tolsma, C. Finn, and D. Sadigh. Learning latent representations to influence multi-agent interaction. arXiv preprint arXiv:2011.06619 2020.
Yuan, X. Weng, Y. Ou, and K. M. Kitani. Agentformer: Agent-aware transformers for pecio-temporal multi-agent forecasting. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 9813-9823, 2021.
Y. Zhu, Z. Wang, J. Merel, A. Rusu, T. Erez, S. Cabi, S. Tunyasuvunakool, J. Kramár, R. Hadsell, N. de Freitas, al. Reinforcement and imitation learning for diverse visuomotor skills. arXiv preprint arXiv: 1802. 00200954, 2018.
14

---

Under review conference paper at ICLR 2023
A PROOFS
We restate the theorem and corollary for convenience.
Theorem 1.
We assume the model po po(â)s, 8, a) [ s) G.)d is achieving optimal re- with the marginal encoder po = Epp [eo(ğ] a, 8)] prior distribution. We can say for the training distribution P(s. a. 9.) = (s a eo (..ss.o0 a) and and testing distribution P(s, a. 9p.) = Pp(s)po To 9,)): If (A [(S, A) and (AA, G.) H(Å, Gp), then I(S, Å) < A).
construction loss Lree 0 on Po(s,a), a). The test policy is po(â)s) = po [ (o 0p)dg
To
Corollary 1. If H A G.) = 0, the assumption H (A, G.) H [A, G,) becomes unnecessary in theorem and have [5 Å) 0 < I(S,A).
A.1 PRELIMINARIES
We denote by ((X) the entropy, by (X the conditional entropy, by YX.Y) the mutual information and by [(X, Z) the conditional mutual information between random variables. Furthermore, the proof relying on the interaction information r(X, Y, 2), an extension of mutual information to three variables. Importantly, the interaction information can be positive or negative. A positive interaction information indicates that one variable explains some of the correlation between the other two while a negative interaction information indicates that one variable enhances their correlation.
Our model @ aa, Te (á) s, 9e) is trained on the dataset Pb(s, a). To achieve minimal reconstruction loss, the model required to predict à = a with certainty, implying H (AAS,, G.) = H HÁASS, G,) = 0.
At test time, latents drawn from the prior (o)) which assume matches the marginal distribution of the encoder, i.e. Ep, (co (9.. a, s)], perfectly.
We use the following equalities:
(X, = I(X,Y) [(X, Z) (and permutations as [(X, 2) is symmetric) (X,Y)Z)) = H(XIZ) - (X Z) (and permutations)
H(X|Y, Z) s (X
H(X,Y) H(X) HY (and permutations)
I(X,Y)
A.2 PROOF OF THEOREM
During training on the dataset Po (s, a) the interaction information is positive because H(A < ((S.A)
T(A,G. S) = T(S,((S A) IS,, I(S, A) (({
H(A(G.) + HHAIS, G.) 0.
(5)
=0
On the other hand, during testing, have (GG, S) 0 because now G, is drawn independently of S from Pe G,). Consequently, the interaction information becomes negative:
1(Å, G,. S) I(G,, S) I S- 50
(6)
=0
15

---

Under review conference paper at ICLR 2023
We also have, similarly to eq. (
G() -H(A)G,)
implying
- 0 [(S A) -H(AG.))
(7)
and hence, if ((A (A(G,)) this gives us the desired result
I(S, A) < I(S,A)
(8)
By assumption, we have
H(G.) -(A +
(9)
Furthermore, because the marginals of G. and G, are matched, we have (G.)) = =G(( and hence the required H (AlG.) = (A(G,) for eq. (8) to hold.
Can we remove H(A, G.) = (Á, G,)) as an assumption? Unfortunately only if (A(G.) = 0, in which case is automatically true (see next subsection). Otherwise this assumption is needed to make sure that the entropy in the system remains comparable between training and testing.
If (HA + (Á, Gp). the main result ((S, A) < I(S, could still hold, but one could also construct environments and encoders in which it does not. The reason is that H(A depends on the distribution p()lg) which changes between training. where is p(slû.), and testing, where it is (D(s) due to the independent drawing of 9.. This can be used to construct environments and encoders that change (AA G) and H (A, G) arbitrarily between training and testing, hence making comparing the mutual information I(S,A) A) and I(S, A) meaningless.
A.3 PROOF COROLLARY
We have p(a g) = (asl)))a((a)s,p.))))) gdd.. We also know that H(ASS, G.).G.) HAIS G,) = 0 and hence To Tā 9) E (0, 1]. Furthermore, if H(AlG.) O, the action is fully determined by G., i.e. To ( 9) To E 1). Hence, because Ja p(s I.)ds- = -S, Pb(8)ds = 1, the switch from p(s (.))
(()) (s)s) does not impact p(alg). so have ([AAG.) ((A
The result that (S, A) 0 follows directly from eq. and (SS, Á) > 0.
B ADDITIONAL RELATED WORK
Unlike our work, agent modelling (Grover al., 2018: [aapoudakis and Albrecht 2020) often assumes knowledge of agent identities in multi-agent systems and aims at learning a useful representation for each identity. In contrast, we neither know the true type g of the imitated agent, nor the identity of external stochastic noise source Ç. Furthermore, applications of opponent modelling in RL settings (e.g., Papoudakis and Albrecht 2020 He et al.) 2016; Raileanu et al.) 2018 Hernandez- Leal et al., 2019 Xie et al., (0200) are generally unconcerned about distributional realism and do not consider distribution shifts.
Behaviour Prediction (BP) also forecasts future trajectorics. Unless future steps are predicted independently of the evolution of the scene (e.g., not auto-regressively) (Chai al. 2019 Cui al. 2019 Phan-Minn et al.
2020
Liu et al. 2021), these methods also suffer from covariate shift in the state visitations (Bengio et al. Lamb et al. 2016). Furthermore, if hierarchical methods are used to capture the multimodality in the
2015
16

---

Under review conference paper at ICLR 2023
data (Tang and Salakhutdinov 2019 Casas al. 2020 vanovvcc and Pavone 2019 Salzmann et al. 2020
Yuan et al. 2021; Hu et al.] 2021), they are vulnerable to the same marginal and conditional type shifts we consider. W hile none of these works take these challenges into account, they often use small discrete latent spaces (e.g., Tang and Salakhutdinov 2019; Ivanovic and Pavone 2019; Salzmann et al. 2020), mitigating the severity of the distribution shifts and future information leakage by limiting the information bandwidth latent types. Furthermore, prediction quality metrics such as displacement-basee less sensitive to yield lower performance due to covariate shift, which primarily impacts interactions with the environment, such as collisions.
metrics or log-likelihood are
As discussed section 3 the conditional type shift is exacerbated by causally confused policies relying the latent type for information about environmental noise. Unlike in most literature causal confusion (De Haan et al. 2219). our nuisance variables are hence not part of the current state, but the learned latent state. Distribution shift is induced not through earlier actions but through sampling from the prior instead of the encoder. Prior work on causal confusion typically relies on problem specific regularization (e.g. Wen et al. 2020 Park al. 202I) or has access to an expert or task rewards (e.g. De Haan et al. 2019 Ortega et 2021). Instead, our work relies on generalisation over latent types to generate counterfactual trajectories. This generalisation is enabled by the information bottleneck and results are refined by the adversarial loss.
C LIMITATIONS AND SOCIETAL IMPACT
While RTC notably improves distributional realism (see section 6 it does not achieve it perfectly, especially in the long tail of the data distribution. This has implications for its use, for example in economic simulations to evaluate policy proposals or in driving simulations to evaluate autonomous vehicles, where this limitation has to be taken into account and the simulation results should not trusted unconditionally.
As RTC is application agnostic, the societal impact depends on where is used. Here, we focus agent-based simulations as we anticipate this to create the highest impact. Examples include better policy decisions through economic simulations, safer autonomous vehicles through driving simulations, better AI in games or improved safety precautions for large crowds of people. For other use-cases, e.g., in armed conflicts, the societal impact will depend the intention of the simulation. Furthermore, we stress many use-cases, precautions have to be taken to account for remaining errors in the learned agents.
more, that for
Lastly, depending on the use case, algorithmic bias has to be taken into account if mode-collapse might be prevented more effectively by RTC for certain strata in the population.
D ADDITIONAL EXPERIMENTAL DETAILS
D.1 DETAILS ON DOUBLE GOAL. PROBLEM
The agent observation
$: It., a.-1] []] E 12 with = [It, 9.,, [I Vr,el, i {u,(}
(10)
contains the 2D position of the agent, 8t, as well as marked locations It. Iu.; of the lower and upper goal. Because the current agent position cannot uniquely identify the currently selected goal. the observation also contains the last agent action a,-1 with the simple transition function 8+41 8, a,, The goal locations are randomly sampled at the beginning of each episode. The lower (upper) goal is always located in the lower (upper) half of the z, y plane. Their horizontal and vertical distances from the initial agent position
are uniformly sampled within rectangular bounds x; E [1.8, 2.2] and ly. (x) E [0.3, 0.7]. Each episode has
a fixed horizon of T 30 steps over the course of which each goal moves by | +(0)]1 0.15 in a random direction.
17

---

Under review conference paper at ICLR 2023
Collision offreed Time
Collision
Offrued Time
Symphony
0.25
0.03
Hierarchy NoFT
5.15
InfOPKAIL C
0.20
InfOMEAIL
6.15
5.10
RTC D
0.19
0.45
9.92
0.45
5.00
308500
150000
200632
50990
106000
150850
203000
S2060
100056
150000
200000
56993
180090
150000
200380
Figure 4: Waymo Open Motion Dataser: Performance on the validation set during training. Distributional realism metrics are not shown their evaluation is high variance on the small validation set.
For the first 10 timesteps, the agent randomly resamples the target goal with P(G = gi) = 0.75 to avoid a simple decision boundary along the r-axis in which experts in the lower half-plane always target goal g The expert action is
de [de
0.1 0 0
a₂ 0.1
Ael
where de
A.
and A,
(g. - -t)).
(11)
0.05
The expert approaches the goal faster along the z-axis, hence creating curved path. To avoid over-shooting, the step-size reduces by V as the agent proceeds towards the goal.
All networks use simple MLPs with two latent layers and a latent dimension of 256. To capture their shape, the discriminator acts entire trajectories, aggregating across time using max max-pooling per-timestep embedding. All policies are parameterised as Gaussian mixture models with two modes. RTC uses a continuous bottleneck of size 2 with additional regularization term C (T) 3KL [ea (g. (W, D) regulate the information bandwidth of ge- We use the BC loss for Lee fe)
32 dimensional
log To a 9.).
Batch size is 1024 for training and 10K for evaluation. Results shown in fig. 3 are evaluated every 100 steps and exponentially smoothed with a decay rate of 0.9. The learning rate is 0.01 for BC (lower learning rates performed worse) and 0.004 for both MGAIL and RTC, which were tuned independently for values E [0.02, 0.01, 0.004, 0.002, 0.001]. f = 0.5 was used to split between Benewder and Bprior (no tuning was performed). Lastly, without further tuning, Audv = 1 was used. Training time is about 7h without hardware acceleration.
D.2 DETAILS ON WAYMO OPEN MOTION DATASET
The Waymo Open Motion Dataset (Ettinger et 2021 (published under Apache License 2.0) consists of segments of length 9s sampled at 10H z. The available training and validation splits in the dataset consist of 487K and 49K segments each, which are used for training and testing the agent respectively. Due to memory constraints, we filter for segments with less than 256 agents and 10K points describing the lane geometry resulting 428K train and 39K test segments. 250 segments from the training split are used for validation select the training checkpoint for evaluation. In each segment, we learn to control two agents at frequency 3. 2, repeating actions three times. Similar to (Igl et al. 2022). the actions of other agents are replayed from the logged data. The collision metric measures the number of segments, in percent, for which at least one pair of bounding boxes overlaps for at least one timestep. The off-road metric similarly detects for how much time the agent's s bounding box overlaps with off-road areas.
18

---

Under review conference paper at ICLR 2023
The state 8; = [s) s (.55), (DS) s{(U)] contains the agent's position and heading { static features ($S) such as lane boundaries, expressed as set of points, dynamical features 8 (DS) such as traffic light states, and the positions and headings of all other road users 8; (RU)
Similarly to section 6.1 we use an additive transition model in which the policy predicts the change in the
agents state (a) s az Dynamic features and roadgraph users are replayed from the logged data,
similar to [ggl et al.
2022
All positions and headings are first normalised to be relative the observing ego-agent. MLPs are used to encode each object and point individually and per-type max- maxxloooing is used to aggregate over variable number of inputs. The resulting three embeddings (one for the ego-agent, one for other road-users and one for the scene), each of size 64, are concatenated and passed either to the policy, discriminator or value function, whose encoders are not shared and which consist of MLPs with two latent layers of size 64 for discriminator and value function and 128 for the policy. The inference encoder ea for RTC only observes future agent positions s(5) which are each concatenated with an eight-dimensional learned positional embedding and individually encoded to dimension 128 and max-pooled along the time dimension. A Gaussian mixture model with modes was used for all policies, although we find empirically that typically only up to three are used after training.
We train for 200K gradient steps and select the model checkpoint for evaluation with the lowest sum of collision rates and off-road time the validation set. To stabilize training for all methods, we discount gradients through time with y = 0.9 and bootstrap from learned value function every 10 steps. We anneal f from I to fmin 0.5 over the course of training. Initially, high values of f encourage meaningful information in ge while lower values address covariate and type shifts and improve performance. A learning rate of 0.0001, which was tuned for MGAIL, was used for all evaluated methods. Each batch contained 24 segments and training was performed on a single V100 (per seed) and required about 4-5 days. We used Aadv = 4.0 and 0.01 for (T) for continuous type representations of size 2. Discrete type representations used three one-hot vectors of size 16, trained using Straight Through gradient estimation. We found performance to be marginally better for three vectors, compared one, without noticeable performance increases for additional or larger vectors. Smaller vectors with only four values only performed slightly worse. The Huber loss Leee uses 30.
D.3 DETAILS ON SYMPHONY BASELINE
Symphony implements the hierarchical policy proposed in Igl et al. (2022) (called 'MGAIL+H' in their results). Agent types represent high-level driving intent and are expressed as a sequence of road-segments be followed. They are encoded into latent vector by expressing them as fixed-length sequence of points { N Each point concatenated with a positional embedding, then encoded individually, and subsequently max-p max-pooled along the time-dimension. The pooled embedding is provided as additional inputs both the discriminator and the policy.
During training, lane sequences are extracted from the given data trajectory т. The prior (e(ğ), which is used during testing when no ground truth trajectories are available, predicts categorical distribution over all possible sequences of lane-segments which the agent could follow in a scene. To allow for variable number of such sequences, the logits are predicted individually per sequence.
D.4 DETAILS ON INFOGAIL BASELINE
Like RTC, IN Li al. 2017) is general method for learning hierarchical agent from demon- strations. What makes it a suitable baseline is that, like in RTC, the higher level policy captures a distribution over alternative trajectories that can be taken. It does in an unsupervised fashion by introducing an
19

---

Under review conference paper at ICLR 2023
Data Distribution
Learned Distribution
(a)
(b)
(c)
Figure 5: Differences between realism, coverage and distributional realism. The data distribution P(Xp) E ( is
shown in green, blue denotes learned distribution (ox)).)) A(X). (a) Data from leamed distribution is realistic, i.c. supp(Xu, (X2) supp(Xp), but distributionally realistic. (b) The learned distribution achieves coverage but distributional realism: the frequencies of modes are not matched. (c) The learned distribution is distributionally realistic. In practice, the dimensionality of X' is often too high, requiring to mcasure distributional rcalism only in selected features ). Consequently, distributional realism in (x)) does not necessarily imply good realism, i.e. lask performance.
additional reward that incentivizes the policy to produce state-action pairs from which an additionally trained discriminator (also called 'posterior') ) can infer the type on which the policy was conditioned. In other words, it rewards the policy for producing distinct trajectories for different types, where the type is drawn from a fixed prior when generating rollouts.
A crucial difference between InfoMGAIL and RTC is that InfoMGAIL's goal is to disentangle trajectories, but it does not target distributional realism directly. particular, because the prior from which the types are sampled is fixed, it might not even be able to properly capture the true distribution of trajectories. This is especially true for the uniform discrete prior used in the original InfoMGAIL paper, which assumes a uniform distribution over trajectory modes. Furthermore, the additional posterior reward introduces bias, potentially harming task performance. Lastly, because mode collapse is not directly penalized the additional loss (only 'non- non-distinctiiveess of trajectories), it might not improve distributional realism at all.
our experiments, we augment InfoMGAIL in several ways:
We not only try discrete latents, but also continuous ones. For continuous priors we use the same GMM as posterior as we use as prior in RTC.
We additionally provide the posterior with the initial state input. Unlike in the examples used in the InfoMGAII, IL paper, we believe that for more complex WOMD data, the current state insufficient to determine modes.
To make it comparable our setup, we optimise it using Info(M)GAIL, i.c. the posterior score of the true type is added differentiable loss term, not reward for TRPO. The network architecture of the posterior is the same one as we used for our MGAIL discriminator.
We greatly increase the number of latent dimensions. In (Li al. 2017). 2 and 3 dimension were used for the two experiments. We tried d E [3, 10,30, 100]. We also tried A [0.01, 0.03, 0.1, 0.3, 1.0] as regularization strength for the additional loss term.
Lastly, we are also adding a BC term to InfoMCAIL we found this stabilizes training greatly.
In contrast to the original implementation, we are not using pre-training and do not make use of additional shaping rewards.
20

---

Under review conference paper at ICLR 2023
D.5 COVERAGE AND DISTRIBUTIONAL REALISM METRICS
While coverage is easy to achieve on the Double Goal Problem, we measure it on the Waymo Open Motion Dataset (section 6.2 using
1
minADE (T) min T
Σ 8(,, Br,e)
using a fixed number of K 16 rollouts per segment. Intuitively, the more modes are covered by a given agent, the closer one of all K rollouts should be to given trajectory from the dataset, resulting in lower minADE.
We want to measure distributional realism the divergence between the expert distribution Pespen (7) and the predicted distribution Pagen(( (f). However, since the space of possible trajectories is far too large to directly measure JSD (Pagent (f) (Pxpeeu)))))) we extract scalar features h from trajectories and measure the divergence on those features. For the Double Goal Problem, we would like to capture whether the agent is approaching 9t Iu- for which extract h2 sygn(yr), i.e. whether the agent is in the lower upper half of the plane at the last timestep. In the driving domain, we measure progress as the total distance travelled over the 9s segment, i.e., hatyle 0 8T). Measuring this distance a straight line avoids measurement noise through swerving or jittering of the agent. Lastly, to measure high-level intent, i.c. whether the agent prefers
going left, right or straight branching points such intersections, we follow Igl et (2022) and extract as feature hear, i.e., the curvature of the lane segments being followed right after possible branching points in the road.
21


Reviews:
[{"review_id": "yeYX5_K9fc", "reviewer": "Reviewer_MSid", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper focuses on the problem of imitation learning in stochastic environments. The stochasticity in the environment is characterized by the parameter  $\\xi$, that is independent of the expert/agent.\nThe goal of this work is to build imitation agent that achieve distributional realism — match the expert distribution covering all modes of demonstrations while achieving good task performance.\nAs a solution approach, the authors attempt to build on and improve previous works on designing hierarchical policies - conditioning the imitation policy on type of expert demonstration- to achieve the stated goal.\nThe hypothesis is that the demonstrations exhibit multiple modes due to difference in types of expert (where type models expert-specific features e.g. persona, goal, strategy or belief) and such types may be sampled from demonstration of available as labels in the data. The authors argue that existing works in type-conditioned imitation learning (especially when such types are latent and sampled from data) either fail to achieve satisfactory mode coverage or compromise task performance. The authors identify two types of distribution shift in the inferred types - marginal (mismatch between posterior and prior type samples)  and conditional (mismatch arising due to the inability of existing type encoders  to disentangle environment noise and agent specific features) - as the reason for inducing causal confusion in the agent, thereby degrading distribution matching and task performance. The authors propose new objective that aims to mitigates these shifts by allowing to train with types sampled from both posterior and prior (unlike previous methods which only consider posterior types in training) and augmenting the imitation loss with information bottleneck term to address the conditional shift stated above. Empirical investigation on a toy double goal problem and more complex Waymo Open Motion dataset is presented to discern various components of the proposed objective and provide comparisons with representing adversarial IL baselines in terms of task performance and distribution matching.", "summary_of_the_review": "While the paper deals with an important problem of imitation learning in stochastic environments, the proposed approach is highly incremental and adhoc. The identification of marginal and conditional distributional shifts appears as a reinterpretation rather than motivations grounded in theoretical or empirical evidence. This coupled with multiple major shortcomings in the empirical exercise makes this contribution below par in its current form and informs my current assessment. ", "strength_and_weaknesses": "Strengths:\n-----------\n+ Imitation learning in stochastic environments is an important problem as most real-world environments have stochastic elements and an approach that can disentangle stochasticity in environment from stochasticity in agent behavior is very useful direction. \n+ The proposed method show consistent improvement in performance across different metrics over considered baselines.\n+ Effective evaluation of imitation learning method is an open issue and the considered metrics, especially Average displacement error and JSD based metrics are interesting and insightful to use for evaluation.\n\nWeaknesses:\n---------------\n- The proposed approach heavily relies on the insight that existing type conditioned approach leads to causally confused agents as the design of such methods require the agents to infer the stochasticity of environment from the latent type encodings whose primary purpose is to provide agent-specific information. However, this entire premise is not given a rigorous treatment and presented very loosely with anecdotal examples without any theoretical or empirical  grounding. This makes the identification of conditional distribution shift as a reinterpretation of previous works in light of the proposed approach rather than a concrete motivation. For instance, the example of pass/yield is used to claim that the policy stops using current observation to infer stochasticity and only relies on latent type. Without establishing theoretically or empirically, that existing strong baselines indeed end up doing this, such a claim is not substantiated.\n- The main contribution of this work is the objective function Eq 2. that is proposed to mitigate the marginal and the conditional distribution shifts. The reconstruction and adversarial loss have been extensively used by previous imitation learning works [1] even for hierarchical type-conditioned policies [2]. Effectiveness of Information bottleneck term has also been studied previously in the context of imitation learning [3]. Further, the explanation of why this term, while helpful in gaining most of empirical performance (see later), contributes to the disentanglement of $\\xi$ and type is unclear and adhoc. The use of sampled prior for adversarial term is novel but it is highly incremental contribution given the existence of previous works on type-conditioned imitation learning. Again, it is not clear why introducing this term in effect, mitigates the marginal shift. \n- There are several strong works that are either missing citations [4.5,6] or discussed in a very hand-wavy manner without providing comparisons with them [2,6]. Further, a great weakness of the paper stems from the unsubstantiated claims by the authors claim that hierarchical methods such as [2,6,7] give rise to conditional and marginal shifts . It is imperative that such statements be backed up with rigorous analysis of those methods to effectively showcase the existence of identified issues.\n- Due to the above limitations, the burden of the work lies on empirical support, however, the paper has several shortcomings in this area:\n    - Lack of comparisons with strong baselines [2,4,5,6,7] is a big miss. It is imperative to show the efficacy of approach as compared to these works as they can be considered as different variants of this approach. \n    - The authors mention that an ablation Hierarchy-NoPT which considers no prior training is similar to existing hierarchical methods but it does not cover the methods cited here. Further, if one considers Hierarchy-NoPT as representative of existing methods, then the improvement with prior training is marginal and hence full approach has very limited contribution. \n    - Minor follow-up: Given the high performance on  minADE for Hierarchy-NoPT in WOMD dataset, the high collision rate seems like a bug and the reasoning for this is not adequate and seems adhoc. \n    - On the double -goal problem, the hyper-parameter for IB is shown to be very sensitive which is a limiting factor as it may not be easy to tune this hyper-parameter. Further, given that this term affects the performance more than any other objective and it is very sensitive, comparisons of the form baselines + IB becomes very useful. For instance, will [2] + IB provide superior performance than proposed approach?\n    - Ablation is only provided on toy task however, there is a wide gap in the characteristics of double goal problem and WOMD and hence an ablation on WOMD is also required. Plus, further granular ablation is needed for ex. does reconstruction + IB provide as good performance?\n    - It is mentioned that learning rate was only tuned on MGAIL to 0.0001 and then used for all methods. This is a strange design choice and tuning should be done specific to each method.  \n    - While there is performance improvement across different metrics, it is very difficult to establish that this gain has anything to do with the mitigation of conditional and marginal distribution shifts induced by other methods, which fails to support the main claims of the paper.\n\n[1] Imitation Learning: Progress, Taxonomies and Opportunities, Zheng et. al. 2021.\n\n[2] Robust Imitation of Diverse Behaviors,  Wang et. al. Neurips 2017.\n\n[3] Variational Discriminator Bottleneck: Improving Imitation Learning, Inverse RL, and GANs by Constraining Information Flow, Peng et. al. ICLR 2019.\n\n[4] Imitation Learning from Visual Data with multiple intentions, Tamar et. al. ICLR 2018\n\n[5] Triple-GAIL A multi modal imitation learning framework with Generative Adversarial Nets, Fei et. al. 2020\n\n[6] Multi-modal Imitation Learning from Unstructured Demonstrations using Generative Adversarial Nets, Hausman et. al. Neurips 2017\n\n[7] Learning a Multi-modal policy via Imitating Demonstrations with Mixed Behaviors, Hsiao et. al. Neurips Workshop 2018", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "y-d5pAKGrx5", "reviewer": "Reviewer_4XDc", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper aims to address the issue of replicating an expert's entire distribution of behaviors in imitation learning. In stochastic environments, an agent's behavior distribution can be caused due to its internal strategies or environment stochasticity. The paper claims that prior approaches suffer due to either not being distributional or because they do not disentangle environment stochasticity. This paper proposes existing tricks, such as using variational models, sampling from a prior distribution while training, and adding an information bottleneck to solve these issues.", "summary_of_the_review": "To the best of my understanding, lack of clarity in writing, unjustified claims, and supposed lack of novelty are my reasons for recommending rejection.", "strength_and_weaknesses": "## Strengths\n- The overarching goal of replicating the entire distribution of an expert's behaviors in stochastic environments is well known to be important.\n- Experiments are reported on a real-world dataset with good practical impact.\n\n## Weaknesses\n- **Writing and lack of clarity**: The paper is written in a way that is very hard to follow and understand the primary contributions of this work.\n    + Many uncommon terms are used for existing alternate terminology, e.g., agent type (for trajectory encoding), hierarchical architecture (for variational methods), and distributional realism (which is already a goal of variational methods).\n    + Lack of structure and flow across paragraphs makes it hard to follow.\n        * The introduction lists many aspects and solutions but does not clearly list what the exact problem being solved in this paper is and how their contributions solve them.\n        * Likewise, Section 3 is quite hard to follow. Is the paper solving two problems, marginal and conditional type shifts, or causal confusion as well?\n    + The assumption of a known differentiable environment suddenly appears on Page 3 without any prior mention or justification.\n    + What are the contributions of this paper's Robust Type Conditioning (RTC) method? From the writing, it seems that all 4 proposed losses are novel contributions. But that is undoubtedly not the case, as all of these losses are well-known and commonly used losses. The writing should be clear enough to distinguish what exists in the prior work and what is contributed by this method. This is also important for knowing what ablations should be performed. More on this is below.\n- **Problem Formulation (Section 3)**:\n    + A lot of issues are there with clarity.\n        * If the majority of the problems in the paper come because of stochasticity from the environment in the input trajectory, then why is Figure 2 explained using a trajectory that has just one (s, a) pair? It makes it very hard to understand conditional-type shifts.\n        * If marginal type shift is an existing problem in VAEs, called the prior hole problem, how is it solved in the prior work, and what is this paper doing differently?\n- **Approach (Section 4)**:\n    + Unclear why the problems mentioned in the paper arise when using standard approaches. For instance, training a standard state-conditioned (not trajectory-conditioned), variational GAIL policy should be able to resolve the problem of mode collapse, which is claimed to be one of the critical problems in the paper.\n    + The paper says that the reconstruction loss prevents mode collapse. But, the paper also argues that the hierarchical (or variational) model is what prevents the mode collapse. The way it is written, it seems the reconstruction loss is a novelty. However, it is unclear to me how that is.\n    + By minimizing the $L_2$ distance between agent true and predicted states at time $t$, how does the policy learn to imitate actions?\n    + Why isn't there any KL term as is common in variational models? Would the $L_{prior}$ be needed if we have the standard KL loss?\n    + The paper claims that $L_{adv}$ eliminates the marginal and conditional type shifts but does not explain how. Two paragraphs later, the paper claims that the conditional shift is not resolved, and thus information bottleneck is needed. Which one is right?\n    + The claimed information bottleneck is not a novel solution and is well-known in the literature.\n- **Experiments (Section 6)**:\n    + Overall, because the paper's contributions are unclear, it is hard to understand what the ablations should be. In the current way of writing, I would imagine that all 4 losses should be ablated, but only \"No Prior Training\" is ablated.\n    + While I appreciate the existing baselines and environments, it's hard to comment on the correctness without understanding the exact contributions of the paper.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "1: The main claims of the paper are incorrect or not at all supported by theory or empirical results.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "MDrZA0NbHa", "reviewer": "Reviewer_JuGQ", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Mode collapse is a problem in imitating multi-modal trajectories. Hierarchical method use latent code to represent the agent type or context variable which is inferred during the training and later used at test time, and the policy should be conditioned on the latent code. Context variable include internal (intrinsic setting of agents) and external factors (coming from environment)\n\nBeyond Covariate shift, the author considers marginal type shift where the prior distribution of agent type changes, where conditional type shift is the external context variable is dependent on the internal context variable. Adversarial imitation learning optimize the policy via a discriminator to distinguish between true trajectory and learned trajectory in a min max fashion where the policy try to imitate the observed trajectory and the discriminator try to distinguish between them. By dividing the data input to be the policy sampled trajectory and the real observed trajectory, the discriminator classification loss could be used as reward for a reinforcement learning agent. \n\nThe author proposed Robust Type Conditioning by incorporating prior during training, which is summarized into a 4-term loss function evaluated on prior policy and encoded policy.\n\n\nResults are tested on Waymo Open Motion dataset. \n\n", "summary_of_the_review": "Adding latent variables to imitation learning is crucial for mimicking multimodal behaviors, as well as being able to generalize better in terms of distribution shift. \n\nThe author offers a very comprehensive view about tackling the problem with fair amounts of experimental efforts, however, I think it is not very well explained, how each of the individual aspects were tackled, for instance, if one remove one of the loss terms, how much the performance will deteriorate? \nWhat is the theoretical contribution of this paper? Is that learning the prior of agent type? ", "strength_and_weaknesses": "Strength:\n1. The author considered a very comprehensive treatment of distribution shift, while most previous work focuses on covariate shift only. \n\nWeaknesses:\n2. I think the method section (section 4: Robust Type Conditioning) is too short, the different components of the loss are not fully explained, which only last for a bit more than 1 page.  For instance\n\n2.1 Why the reconstruction loss prevents mode collapse by penalizing the agent for being unable to mimic \\pi?\n2.2 Why the adversarial loss could eliminate marginal and conditional shift?\n2.3 Could you explain more about the information bandwidth of \\hat{g}_e?  Why the prior on agent type could filter out the environment (external) noise? ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "9OguoBqXoak", "reviewer": "Reviewer_iTsU", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes RTC, an imitation learning approach for stochastic environments that aims to learn policies that avoid mode collapse and cover a wide range of behaviors shown in demonstrations. The paper argues that current methods relying on inferring agent types to explain multi-modal behaviors in stochastic environments suffer from Conditional Type Shift: the agent latents are encoded using a full trajectory and therefore do not account for environment uncertainty, while at test time, agents do not have access to the future, creating a shift between how the agent was encoded at training time and how it is deployed now. To address this, RTC uses proposes to learn a policy by using both the prior and posterior over agent types during training, making sure that trajectories under both distributions are likely according to a discriminator and that the likelihood of both distirbution is high.  \n", "summary_of_the_review": "The paper addresses an important problem, which is well motivated, and described why it arises. While there are some details (see weakness/details) that I would really like clarified, the approach for RTC addresses some of the limitations of existing IL approaches with latent codes for agent type. The results, both the simple environment to provide more intuition about the method and the Waymo Open Motion show the value of this method. Hence, I think this paper should be accepted.", "strength_and_weaknesses": "**Strengths:**\n\n- Motivation\n\t- This is an important and well motivated problem, there are many scenarios in which having agents that cover multiple modes is important, and most environments are stochastic or become stochastic as multiple agents are added.\n- Clarity\n\t- Section 3 gives a clear and intuitive overview of the conditional type shift problem that this work is addressing. I wish some of it was more present in the introduction. \n\t- Very clear overview of RTC in Figure 3 and generally over setions 3 and 4. I value that authors provide explanations of the effects that every loss is trying to mitigate, and propose strategies for different kinds of environment conditions (such as having vs no having access to environment gradients).\n\t- Very clear analysis of results, for both benchmarks.\n- Novetly\n\t- The work proposes a new approach to encode variations in policies learned via imitation learning. The core idea builds on top of GAIL and other works encoding agent variability using latent codes, but introduce novelty to separate agent behavior variability vs environment variability.\n- Experiments:\n\t- Approach is tested against strong baselines and sound ablations, with a simple environment that is easily interpretable and a more challenging real-world environment.\n\t- SOTA results in the Waymo Open Motion Dataset, compared to strong baselines, in terms of performance and mode coverage. The results, particularly the comparison of Hierarchy and RTC shows that the proposed method allows to both cover a high number of modes, but also deal with covariance shift that arises from using a posteriro during training.\n- Limitations well addressed in conclusion.\n\n**Weaknesses:**\n- More details on claims\n\t- It would be valuable to provide a reference or more details on some of the claims made in the paper, such as the failure to cover modes of adversarial approaches (in intro).\n\t- I am not sure I see why encoding encoding all the stochasticity into the agent is a bad approach. In the example in Figure 1, a model-based approach would need to account for stochasticity of the other agents, but the behavior of the agent we are imitating, which only has information of the current and previous states could be controlled by their internal factors, such as risk tolerance to cross or not. A risk tolerant driver will cross when a car is approching, but will stop if the car has already started crossing, avoiding problems. A non-risk tolerant driver will stop directly. I am not sure why we need to explicitly encode the env stochasticity. \n\t- Isn't it a high assumption that there is a differentiable environment? How well does this method work if having to use reinforcement learning?\n- Experiments:\n\t- For the double-goal environmnet, it would have been interesting to test the effect of environment stochasticity, given that this is one of the claims of the paper. How do the baselines change where there is more/less stochasticity.\n\n**Details**:\n- I see that the Lprior, maximizes the likelihood of \\hat{g}_{e} and \\hat{g}_{p}. Would it make sense to minimize the KL divergence between the 2? Is there a degenerate case where both prior and posterior have high likelihood according to p_theta, but non over-lapping support? What would prevent that?\n\n- Where is the value function in the method proposed here, compared to Symphony? Does it come from the discriminator? Why does it matter if the policy is optimized end-to-end here, instead of using policy gradient?\n\n\n- Caption in Figure 4 could explain more what the curves in the bottom row mean.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject