Robust Imitation of a Few Demonstrations with a Backwards Model
Authors: Jung Yeon Park
Lawson L.S. Wong
Venue: NeurIPS 2022 Accept (2022)
Keywords: imitation learning
model-based imitation learning
behavior cloning
covariate shift
robustness

Abstract: Behavior cloning of expert demonstrations can speed up learning optimal policies in a more sample-efficient way over reinforcement learning. However, the policy cannot extrapolate well to unseen states outside of the demonstration data, creating covariate shift (agent drifting away from demonstrations) and compounding errors. In this work, we tackle this issue by extending the region of attraction around the demonstrations so that the agent can learn how to get back onto the demonstrated trajectories if it veers off-course. We train a generative backwards dynamics model and generate short imagined trajectories from states in the demonstrations. By imitating both demonstrations and these model rollouts, the agent learns the demonstrated paths and how to get back onto these paths. With optimal or near-optimal demonstrations, the learned policy will be both optimal and robust to deviations, with a wider region of attraction. On continuous control domains, we evaluate the robustness when starting from different initial states unseen in the demonstration data. While both our method and other imitation learning baselines can successfully solve the tasks for initial states in the training distribution, our method exhibits considerably more robustness to different initial states.

Robust Imitation of a Few Demonstrations with a Backwards Model
Jung Yeon Park
Lawson L.S. Wong
Khoury College of Computer Sciences Northeastern University
Khoury College of Computer Sciences Northeastern University Boston, MA, USA 1sw0ccs .ee.ccs..ee..c edu
Boston, MA, USA
park. jark
Abstract
Behavior cloning of expert demonstrations can speed up learning optimal policies in a more sample-efficient way over reinforcement learning- However, the policy cannot extrapolate well unseen states outside of the demonstration data, creating covariate shift (agent drifting away from demonstrations) and compounding errors. In this work, we tackle this issuc by extending the region of attraction around the demonstrations so that the agent can learn how to get back onto the demonstrated trajectories if it veers off-course. We train a generative backwards dynamics model and generate short imagined trajectories from states in the demonstrations. By imitating both demonstrations and these model rollouts, the agent learns the demonstrated paths and how to get back onto these paths. With optimal near- optimal demonstrations, the learned policy will be both optimal and robust to deviations, with wider region of attraction. On continuous control domains, we evaluate the robustness when starting from different initial states unseen the demonstration data. While both method and other imitation learning baselines can successfully solve the tasks for initial states in the training distribution, our method exhibits considerably more robustness different initial states.
Introduction
While reinforcement learning (RL) has shown remarkable success in many challenging domains 21 38 36), tasks with sparse rewards and long horizons still remain extremely difficult solve. In such tasks, positive reward is only encountered after the RL agent reaches the goal after a long sequence of actions, meaning that cannot learn any useful signals until this occurs (typically the agent must reach the goal several times to learn reliably as well). Furthermore, the learning signal decreases exponentially with the horizon, which combined with slow gradient-based updates, can lead to catastrophic forgetting even after the agent learns to reach the goal.
These demonstrations can be used in the supervised setting where the agent imitates the expert's behavior, termed imitation learning (IL). However, naive behavior cloning (BC) of the expert's trajectories suffers from covariate shift: the agent's 's policy drifts away from the expert's, s, which leads to compounding errors due to RL's sequential nature. Furthermore. the distribution of states given in the demonstrations often has low-dimensional support with respect to the entire state space. As such, the agent cannot extrapolate correctly when outside of the demonstration data. Many approaches solve this issue have been but such approaches require an interactive expert to query the correct actions.
Expert demonstrations can help RL agents solve diffieult tasks 30 43 22
proposed 33 39
16],
Another way
demonstrations is combine imitation learning with reinforcement learning in form of learning from demonstrations (LfD) 65 13 9. In this case, demonstrations do not simply as supervised labels and can guide the agent's exploration, and also as augmentations to good
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Training
(b) Evaluation
(c) BC
(d) BMIL
The policy is trained on a specified set of initial start and goal states and is evaluated at different start states, (c), (d): BMIL uses generated reverse-time rollouts from backwards model (green arrows originating from the demonstration) to learn wider region of attraction (green) around the demonstration data (orange) than BC.
Figure I: (a). (b): Robustness:
data samples. These LfD approaches either use demonstrations pretrain the policy auxiliary imitation loss in conjunction with the policy update [30 22 or modify the reward function such that the agent is rewarded when imitates the demonstrations 50 However, these methods require interactions with the environment whereas we do not assume such access in our setting.
an
One primary concern of relying on demonstrations is that they costly obtain, especially in real-world applications. Requiring an operator providing corrections in-the-loop to handle covariate shift is often prohibitive well. Given only few offline trajectories demonstrating successful task completion, an agent ought to be able to replicate the behavior from similar starting conditions, even there are small perturbations along the way, and correct itself when necessary. We primarily interested in this setting. In this work, we aim to minimize the number of demonstrations necessary for sparse-reward reward tasks, while preserving successful task completion. To tackle covariate shift, we seek an approach that will be robust, in the sense of Figures Ia and [b] the agent is only trained with few demonstrations from single start state, but evaluation time, must generalize its behavior to variety of unseen start states, If the agent can successfully complete the task from states adjacent to the demonstrated trajectory, then can recover from small deviations from the demonstration.
Inspired by the notion of "funnels" in robotics (0) and feedback control 2 19], we introduce reverse-time generative model that can generate possible paths leading the agent back onto the demonstrations. These reverse rollouts provide useful information because every rollout ends within the support of the demonstration data. Assuming that the demonstrations lead to the goal, imitating these generated reverse rollouts and the original demonstration data allow the agent to reach the goal from more starting conditions, including unseen ones. As illustrated in Figures Ic and [d] typical behavior-de behavior-cloning (BC) methods focus learning on the small number of demonstrated states, whereas our proposed approach, Backwards Model-based Imitation Learning (BMIL), uses the reverse rollouts to learn wider region of attraction around the demonstration. We validate our approach on number of long-horizon, sparse-reward continuous-control tasks. Even from 5-20 demonstrations, BMIL. provides a significant increase in the region of attraction and robustness on many domains compared BC, or to using a forward dynamics model.
We summarize our contributions as follows:
We propose an imitation learning method that pairs backwards dynamics model with a policy and train on both demonstrations and imagined model rollouts.
In the restrictive setting of an offline expert and no access to environment interactions, we show that backwards model can improve robustness over behavior cloning.
variety of long-horizon, sparse-reward domains demonstrate that BMIL can noticeably extend the region of attraction around the demonstration data, even when trained very small subsets of the state space.
Our experiments
2 Related Work
Imitation learning has long history (27 and well-studied, as documented in various surveys
25 [I]. The challenges of covariate shift and compounding errors are also well-known 27 32
Most solutions involve on-policy imitation learning, where environment interaction and interactive querying of the expert allow for the agent's distribution to match that of the expert 33 39 More closely related to our approach are methods that modify or augment the demonstrations to increase
robustness. Laskey et al. [16] inject noise into the supervisor's 's policy during training to force the

---

34 demonstrator to provide corrections. Luo et al. [8] learn a dynamics model from demonstrations to conservatively extrapolate value function that encourages the agent to return the expert data
distribution. Generative approaches have also been used in imitation learning [TO 46 but their focus
not on robustly following few goal-reaching demonstrations.
Time-reversibility has been explored in RL, often as a form of regularization 41 23
49
29
Reverse-time dynamics models, also called predecessor models, have also been used RL 3 6 48 77. However, in all cases, the reverse-time dynamics model is used as either an alternative to the forward-time dynamies model or as an auxiliary model in addition to the standard forward-time dynamics model, in order to mitigate model-compoungg reverse-time dynamics model can accelerate RL and enable greater sample-efficinny. In this work, we take a different perspective where the reverse-ting timse dynamics model is used to generate possible, unseen paths that can lead back to the expert demonstration and thus to the goal, thereby improving robustness in following the demonstration. Our work is also similar to [45, where reverse-time dynamics model is used to generate possible trajectories; however, their focus on offline RI, where the generated trajectories are used to connect distinct sets of states in the offline dataset.
37
15
17
errors. The result is that the
3 Method
3.1 Preliminaries
We model the setting as a Markov Decision Process (MDP) with continuous state space S and continuous action space A. T S x A + S the transition function that defines the distributions of the next state $t+1 E S, given the current state E S and action a E AA taken at timestep t.
The objective of imitation learning is learn a policy no, parameterized by 6, that matches the expert's policy Tg. We assume that the expert generates demonstrations Daemo by rolling out its policy Næ in the environment. Note that consider the more restrictive case of where demonstrations only consist of transitions (s, 8') and not rewards, where s' is the next state.
Behavior cloning (BC) is a form of imitation learning where the policy learns to imitate expert actions via supervised learning. The policy learned by behavior cloning found by minimizing the negative log- likelliddoo over the demonstration data
LBC (e(e [o (a s)]
Note that BC does not require environment interactions and can be considered to be offline. Further- more, expert behavior inferred only from demonstrations without access the expert policy.
Compounding errors in behavior cloning As pointed out in 32 behavior cloning suffers from covariate shift, where errors in the policy can compound and lead the agent states where cannot recover. Intuitively, this occurs as states in the training data Daemo is a small subset of the entire state space S and it is difficult for the policy ma learn the optimal action for states outside of the training data. If during rollout, once the policy me makes an етrоr and leaves the Ddemo it may encounter completely new states, leading to the compounding of errors. Furthermore, as the agent moves farther away from Ddemo- there is very little hope for it to make the correct action and move back onto the distribution of the training data.
3.2 Problem Setting
We consider the same setting as BC, where we do not assume access to the environment or the expert policy during training, and only expert demonstrations without rewards are given. Furthermore, assume that the expert demonstrations are given in the MDP where the sets of initial states So and goals G are both very small subsets of the entire state space. An example of this scenario is maze environment where the agent starts from the same initial state and tries to reach a fixed goal. Formally, define the demonstration trajectories Tdemo coming from probability density
T
P(Tdemo T p(0))
II
TE T a;)p(((411 sc. at):
t=
where 80 E Sq, and sT G. our experiments, So and G consist of single start or goal state and/or the e-ball of its neighborhood. As such, only a few demonstrations required to learn a stable
3

---

optimal policy. Note that this is different from domains in previous work [30 4 which consider random goal states, requiring many more demonstrations to learn optimal policies. We also assume that the expert policy is optimal in the sense that all demonstrations successfully reach goal. While this is not strictly necessary in our method, our setting does not include rewards in Tdemo and thus we cannot discern whether demonstrations are optimal. This allows us to ignore the issue of modifying rewards as done several offline RL algorithms [5] [4]. In order to use task completion success rates as evaluation metric in our experiments, we consider only optimal demonstrations.
Our objective is to learn policy that robust to policy errors when imitating expert behavior and can learn to reach the goal from variety of initial states. This is different than the multi-gos or multi-task setting, where the agent learns to solve multiple goals or tasks, usually from small number of initial states. More formally, the robustness of the policy R(Ta) is defined as
R( E [1 T, E 9)),
where the expectation is taken over SR. where SR is strict superset of Sg. Note that our measure of robustness is somewhat coarse, in that we do not consider the shortest path to reach the goal from every start state (which would probably require more information such as diverse trajectories or environment interactions). Instead, we seek to extend the region of attraction around Demoo such that the learned policy can still reach the goal.
As we consider continuous states in our work, we measure the robustness R using samples from SR- where we randomize either some or all of the state dimensions. For example, in robotic manipulation domains, we vary the position of the gripper we are primarily concerned with being able to learn robustness from a variety of different starting positions. other domains, we are interested in the policy's s ability to recover from arbitrary initial states and so vary not only the agent's 's starting position, but also the initial joint positions and velocities by adding uniformly random noise.
Throughout, we assume that there exists an action E A that allows the policy to go towards D Demmo when a state s of Ddemo- This scenario true for many navigation and physics- based domains we ignore rare circumstances such as irrecoverable unsafe states or the breakdown of the agent. We exclude such cases and assume that state transitions are reversible. We discuss some possible ways to incorporate irrecoverable states in Section 6
3.3 Backwards Model-based Imitation Learning
In work, we use backwards dynamics model provide more synthetic training data to the policy and therefore increase the policy's 's robustness. We call our method backwards model-based imitation learning BMIL.
Backwards model The backwards model is a probabilistic generative model defined as B = p( St, az 8141). This model estimates the conditional distribution of the reverse time dynamics. takes in the next state and outputs the previous state and previous action. As we consider only continuous state and action spaces, we implement B as a conditional Gaussian, parameterized by ø.
The backwards model decomposed into two functions an action generator and previous state generator. The action generator BA predicts which action was taken order to land in the next state. There may several such actions from different states that can lead to the next state. Thus the action generator implicitly encodes backwards policy. It is important for this backwards policy to closely match the learned forward policy x but be different enough to generate diverse new rollouts for the policy train on. The previous state generator Bs predicts the previous state given the next state and previous action taken. The goal of this generator is to accurately predict the backwards dynamics.
Bs-Bs = p(a: I (541))p((sr de; 8t+1),
As we consider the setting with no access to the expert or the environment, B is trained only on Ddemo. The action generator and previous state generator are jointly trained by maximum likelihood
LB =
t=u log p(a; s(41) + log p(s; ag $141);
(1)
where 8t, de, S141 is the state, action, and next state, respectively, at timestep t.
Model rollouts Given expert demonstrations Tdemas we use the backwards model to generate possible several short reverse rollouts or traces TB, starting from every state in Tdemo- As all of these

---

Algorithm Backwards Model-based Imitation Learning (BMIL)
1: for N epochs do
2: Train backwards model parameters & using Eqn. ①
3: Generate K model traces тB from every state E D demo and store them in Dar-
4: for . steps do
5: Sample mini-batch of (s, a) from Ddemo and Das at a fixed ratio.
6: Update policy parameters 0 using Eqn. (2).
7: end for
8: end for
traces end on states within the demonstration data 8 E Ddemo and as all demonstrations reach the goal, following these traces will eventually lead to the goal. For all $1...., 81 in Tdemo: we generate K' traces ТB in a time-reversed manner, where we start from the last state action pair ($H aH) E D Ddemo and then predict (st, ae) for timesteps As we assumed that there are no irrecoverable states in our setting, the rollouts reflect possible paths that the agent could have taken reach 8 Ddemo- If the reverse time model B is accurate and the previous action generator BA gives sufficiently diverse actions, the traces тв are then samples from the region of attraction "funnels" around every state along the demonstration. As we assume all demonstrations reach the goal, these samples from the funnels can be used to learn a robust policy Te. as it can follow the traces onto the optimal path.
H ...... 1. These traces collected into a buffer Dar
Action selection strategy for BA As the backwards model B is trained only on a limited number of expert demonstrations, it is likely that B can only learn accurate reverse-time dynamics for states contained within or close to the demonstration data D Demmo
47 Thus repeatedly rolling out BA would only generate traces whose state-action pairs are contained within Demoo and would not help with learning robust policies. However, we would like to generate diverse traces with new unseen state-action pairs in order to robustify the policy. To balance the model misprediction accuracy with generating plausible state-action pairs, we perturb only the first action generated from BA and not the subsequent actions and also use short horizon lengths for the traces. Note that we are essentially choosing good action selection strategy for Ba. Let ag be the action that the expert would take. A good action selection strategy would place more probability mass closer to the support of the Ddemo providing a "cover" of p(ag (s) but with a wider tail to provide diverse rollouts. As our backwards model B is probabilistic (implemented as a conditional Gaussian with diagonal covariance). we can easily perturb p(agls) (s) by increasing the distribution's 's variance.
Let (N, o²) be the previous action output by BA. We consider two ways to generate action 1) simple scaling of a and 2) resampling a new action a' by adding uniform noise, a' = a + k, k = u[-b, b), where b is fixed hyperparameter. For the scaling strategy, we further scale by the entropy of the probability density as we wish to make the distribution "wider" for peaker distributions.
Algorithm Our method BMIL is outlined in Algorithm T Given expert demonstrations with tuples of the form (st, 0z, $141), we train the backwards model using Eqn. [ to estimate the reverse-time dynamics p(st, 0g st+1). We train our policy me on both the demonstration data Ddemo and the model traces TB by sampling from both at a fixed ratio and using maximum likelihood,
(2) where pa is the probability of sampling from Datemo- As our aim is to learn a robust policy while still succeeding at the original start states and goals, we sample from the demonstrations at a higher ratio than the model traces. Note that BMIL does not depend on the type of imitation learning policy. Any algorithm can be used as long as the policy can be trained with samples of the form as.
L = PaC + (1 (1-Pa)En (s.a)~Tg |- log (a })],
4 Experiment Design
4.1 Environments
We validate our approach on several continuous control domains: 1) the Fetch robotics environment 26). 2) maze navigation with two different agents, and 3) Adroit hand manipulation (30). Figure 2 shows sample images of some environments. For the Fetch robotics environments, we consider the
5

---

Betch-PickAndPlace
MlzzeAnnCorrido7777
Adroit-Relocat Relocate
Figure 2: Sample images of some considered environments.
"Push" and "PickAndPlacee tasks where the objective is to control a Fetch end effector to either push an object to the goal pick object and place the target location. For the environments, we consider 3 mazes of increasing difficulty, where an agent must learn to move itself and then reach the goal. We use both simple point and 29-DoF ant agent. For the Adroit environment, the 'Relocate" task, where one must control a 24-DoF Adroit hand to pick up a ball and move target location. All domains the MuloCo simulator 42 for total of 9 distinct domains. All environments have sparse reward structures, where either every step has a constant negative reward until the goal reached (Fetch) or only the goal has non-zero reward (Maze, Adroit). In particular, the Maze and Adroit environments are quite challenging as they both consist of controlling the agent's joints perform locomotion (maze) or dexterous manipulation (Adroit) over long horizon. More detailed descriptions of each environment including its observation space are provided in Appendix[A]
4.2 Demonstrations and Implementation Details
To generate demonstrations in the Fetch and Maze domains, we train an expert policy by adding the goal position to the state, as in goal-oriented learning, and use off-the-shelf RL algorithms (28, g) For the Adroit domain, we a pre-trained policy from (30). We use 5 demonstrations on the Push task and 10 on the PickAndPlace task, and 20 demonstrations for all Maze and Adroit environments.
For the policy, we use neural networks with 3 fully connected hidden layers with 256 neurons and ReLU activations. For the backwards model, we use 4-layer MLPs with 256 hidden units for both the action predictor B and previous state predictor Bg and use diagonal Gaussian distributions. To train the policy, we use pd 0.5 for the Fetch environments, Pa = 0.8 0.95 for the Maze environments, and Pa = 0.8 for the Adroit environments. We find that higher ratios are necessary for longer-horizon and more complex domains. For the model rollouts, we use the variance scaling action selection strategy for the first action only and use increasing rollout lengths for all domains, similar [12. For a more detailed discussion of experiment details, see Appendix B Our code for the modified environments, generating expert policies, and running all experiments are available at https: htt///gitttu thub com/ wowmjjyark//bblt//
4.3
Evaluation
We evaluate BMIL against behavior cloning (BC) and VINS (18]. VINS specifically aims to learn value functions robust to perturbations using negative sampling and the induced policy learns self- corrective behavior. VINS chosen it is most relevant
setting; other methods such as DART 16 or SQIL 31 require either an online expert or environment interactions.
We the same number of demonstrations for all methods and also keep the same policy network architecture and the total number of policy gradient steps equal across all methods. We train both the policy and backwards model until the backwards model loss converges. Note that goal is not to solve the training task faster but rather to robustify the policy using the backwards model. Additionally, wish solve the task various starting conditions while still being able succeed at the original initial start-goal states.
To evaluate the robustness of the learned policy, we vary the initial states and compute task success rates. For Fetch, we fix the initial gripper, object, and goal position during training and vary the gripper's x and position within the table boundaries during evaluation. We use 10,000 samples

---

Robustness (%)
Relative BC
BC
VINS
BMIL
BC
VINS
BMIL
Push (5 demos)
12.1
12.8 10.4
14,6.00 10.6
1
1.06
1.21
FEtCh
PickAndPlace (10 demos)
4.1
3.4
17.5 10.9
0.84
4.31
UMaze
49.0
47.8 13.5
1
0.81
0.98
Point
Room5x11
36.8.3.4
17.3
38.613.4
0.47
1.05
(20 demos)
Corridor7x7
33.7=15
37.7
38.9
1.12
1.16
MAZE
UMaze
63.0
24.7=2.1
64.8
0.71
1.03
Ant (20 demos)
Room5x11
33.2
30.2
29.1=00
1
0.91
0.87
Corridor7x7
21.7=
19.6
17.6=005
0.90
0.81
ADROIT
Relocate (20 demos)
7.9
3.8
13.3
0.48
1.68
Table I: Robustness evaluation for Fetch, Maze, and Adroit environments over 400, 100, 100 trials, respectively. The bounds indicate 95% confidence intervals. BMIL improves robustness considerably over BC in most environments.
during evaluation. For Maze, we initialize the agent to random start position within discretized grid of the maze and also add random uniform noise to the agent joints' qpos, We sample 100 initial states for each discrete grid cell and compute the success rate. Sampling points for each grid cell gives us an idea of which positions are easy for the agent to reach the goal. Intuitively, such positions would be those near the goal and the demonstrated path. For Adroit, generate 1,000 random initial states by adding uniform noise to the qpos of the hand.
5
Results
Our experiments aim to answer the following questions: 1) how robust of a policy does BMIL. lear?? and 2) what components of BMIL are important improve robustness?
5.1 Robustness evaluation
The robustness results are shown in Table We note that the absolute robustness percentages are generally low for all methods because of the difficulty in extrapolating from limited demonstrations (5 20) with single pair of initial start and goal states. We therefore also include the relative improvement over BC.
In the Fetch environments, BMIL substantially increases robustness over BC and VINS. In particular, our method has an approximately 20% and 330% improvement over BC on Push and PickAndPlace, respectively. VINS on the other hand performs similarly to BC. We see similar pattern on the harder Adroit environment, where BMIL improves robustness over BC by 68%. For the Maze environments, BMIL generally outperforms BC for the Point agent, while the robustness is decreased for the Ant agent. Somewhat surprisingly, BC performs quite well on the long-horize Maze domains. may be that BC has some built-in extrapolation capabilities or that the backwards model may need better latent representations with more powerful networks.
Empirically, we can see that having short reverse rollouts from the backwards model and using only slight perturbations still helps to increase robustness, even though the traces contain some model misprediction errors. We hypothesize that these traces do not necessarily need to be accurate in order benefit the policy and simply need be plausible paths that lead to the demonstrations. may be that having the general correct direction contained in the traces is sufficient for the policy to eventually reach states the demonstration data.
The success rates during training are shown in Table in Appendix C.I BMIL. achieves success rates close 1 for most domains, suggesting that increased robustness does not necessarily come the cost of decreased performance on demonstrated start-goal states. On the other hand, VINS cannot consistently succeed during training, though its robustness is similar to BC.
Visualization of robustness Figure shows which starting positions succeed during the robustness evaluation for Fetch. The green points correspond to successful episodes. We can see that both BC and BMIL succeed more frequently when starting from nearby the demonstration data (approximately
7

---

BC
BMII.
BC
BMII.
(a) Push
(b) PickAndPlace
Figure 3: Visualization of robustness on Fetch for BC and BMIL. We vary the gripper's æ, position evaluate robustness. The green and gray points denote successful and unsuccessful initial positions. The red, orange, and blue points denote the initial start, object, and goal during training. BMIL learns much larger region of attraction along the path from the initial start (red) to the goal position (blue) and even succeeds in some other areas much farther away.
BC
BMIL
BC
BMIL
(a) PointRoom5x!1
(b) PointCorridorxx7
Figure 4: Visualization of robustness on PointRoom5x1 and PointCorrido77x7 for BC and BMIL. The maze discretized into grid and 100 random states sampled from each grid positions (cach random state also adds noise to the agent's initial joint positions/velocities)) Bright yellow corresponds 100% success rate while dark purple corresponds 0%. The regions of attraction for both BC and BMIL are similar but BMIL succeeds more often within its region of attraction.
straight line from the start (red) to goal (blue)). However, larger region of attraction over BC and even succeeds at points that are much farther away from Daemo- We hypothesize that instead of perturbing a single state within Ddemo as done in VINS, learning short reverse rollout from this state allows BMIL to learn optimal paths from states much farther away from Deemo leading higher robustness values. Figure 4 shows similar visualization for some Point maze environments where the agent's s position discretized into grid. BMIL learss region of attraction that is either slightly bigger similar in size to that of BC but has higher rate of success at each cell.
see that BMIL learns much
5.2 Additional Experiments
Forward Backward Dynamics We first analyze the utility of backwards forward dynamics model. On the Fetch environments, we train BMIL with a forward dynamics model p(s' s, a) and compare against the original backwards model p(s, s'). The forward model is implemented nearly identically to other n-step model-basee RL algorithms (e.g. (12), with the exception of no environment interactions. As with the backwards model, generate rollouts from the forwards model starting from demonstrated states and train the policy on both the demonstrations and traces.
Robustness (%)
Relative BC
BMIL (Forwards)
BMIL (Backwards)
BMIL (Forwards)
BMIL (Backwards)
BC
BC
Push (5 demos)
12.1
12.4+0.6
14.6=0.6
1.03
1.21
PickAndPlace (10 demos)
4.1
4.1
17.5
1.03
4.31
Table 2: Forwards (p(s' 8, a)) Backwards (p(s,a a 8')) dynamics model: The forwards dynamics model performs similarly to BC and does not increase robustness.

---

0.25
Rabusiness
0.20-
Pabustness
0.2
method
bom
0.10-
0.1
0.05
0.00
0.0
None
RS(D.1)
RS(0.3)
R5(0.5)
Sch10
SCI30)
sc(sa)
1
10
15
20
Action Selection Method
No. demonstrations
(a) Action selection strategy
(b) No. of demonstrations
Figure (a) Action selection strategy: effect no perturbation (None), resampling (RS), and scaling (SC) robustness on PickAndPlace. The numbers in the brackets indicate coefficients. We see that perturbing the first action is beneficial compared to no perturbation (None) method and allows backwards model to generate diverse traces. (b) Number of demonstrations PickAndPlas more
demonstrations increase robustness for BC and BMIL, but BC plateaus at much lower level.
To generate model rollouts from demonstration states, The total number of parameters is kept approximately constant for the forwards and backwards models. As shown in Table the forwards model offers little to benefit over BC for both Push and PickAndPlace, suggesting that the backwards model is required to produce a robust policy.
the action from the policy no (s).
Action selection strategy We different action selection strategies in trace generation in Fig- ure 5a (and Figure [0 in Appendix C.2). Compared to no perturbation (None), we see that either action selection strategy improves robustness it can lead more diverse trajectories unseen within the support of the demonstrations. We use the variance scaling strategy SC(30) for all Fetch experiments
more stable than SC(50).
Number of demonstrations We also study our method's performance with varying numbers of demonstrations. As shown in Figure [5b] both BC and BMIL improve in robustness with more demonstrations. but BC plateaus at much lower level. On the other hand, BMIL requires slightly more (10) demonstrations than needed (3 demonstrations are sufficient for BC to succeed during training) in order to train the backwards model (Figure I in Appendix
Computation budget As BMIL trains both the policy and the backwards model, it requires more total gradient updates than BC. On Fetch domains, BMIL uses approximately 5x more computation than BC. We train BC for more steps to match or exceed BMIL's computation budget, as shown in Fi is given 1x 20x computation budget). However, more training for BC does not improve robustness and seems to have a harmful effect on robustness.
Training model first and then the policy As offline method, BMIL does require the backwards model to be trained in a single loop along with the policy. We can first train the model first and then train the policy. We compare training the model first and then the policy with the process outlined in Algorithm [ We find there are no noticeable differences when using this model first approach on the Fetch domains, seen Table 6
0.15
0.15
Rabustress
0.10
Robustress
0.05
0.05
0.00
Boloxt
BC120x)
Bclsxt
BC120x)
Push
PickAndPlace
Figure 6: Computation budget: we increase the number of gradient steps for BC by 1 - 20x for Fetch. BMIL has roughly times more total gradient steps than BC(Ix) due to the backwards model update. More BC gradient steps do not increase robustness.

---

Discussion
This work proposes method tackle the issue of covariate shift in imitation learning. We consider the restrictive setting where the expert is offline, where its behavior can only be inferred from demonstrations, and no access additional environment interactions. Specifically, we show that pairing generative backwards model with behavior cloning can allow policy to learn a wider region of attraction around the demonstration data. By rolling out imagined traces from states within the demonstration and perturbing actions to generate diverse traces, BMI.L learns wider funnel than naive BC. Through experiments on several oon--horizon, sparse-reward. continuous control domains, BMIL noticeably improves robustness when trained on narrow set of initial start and goal states and evaluated random starting positions.
There are many possible extensions for future work. BMIL does not necessarily preclude the use of image observations
only assume that slightly perturbing action will lead to next states close to the original next state. However, to handle images, our approach likely requires an additional encoder and possibly more complex network architectures and augmentation techniques. Another interesting avenue could be to quantify how an increasing coverage of state space contained within the demonstration data affects robustness for both BC and BMIL. Finally, one could consider the setting of irrecoverable states and either resample rollouts containing such unsafe states or incorporate a measure of safety within the backwards model when generating model rollouts.
Acknowledgments and Disclosure of Funding
This material is based upon work supported by the National Science Foundation under Grant No. 2107256. This work was completed in part using the Discovery cluster, supported by Northeastern University's 's Research Computing team.
References
[1] Brenna D. Argall, Sonia Chernova, Manuela Veloso, and Brett Browning. survey of robot learning from demonstration. Roboties and Autonomous Systems, 57(5):469 483, 2009. ISSN 0921-8890.
[2] Robert R. Burridge, Alfred A. Rizzi, and Daniel E. Koditschek. Sequential composition of dynamically dexterous robot behaviors. The International Journal of Robotics Research, 18(6): 534-555, 1999.
[3] Ashley D. Edwards, Laura Downs, and James C. Davidson. Forward-backwaar reinforcement learning, 2018
[4] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine, D4rl: Datasets for deep data-driven reinforcement learning, 2020.
[5) Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International Conference Machine Learning, pages 2062- 2019.
[6] Anirudh Goyal, Philemon Brakel, William Fedus, Soumye Singhal, Timothy Lillicrap, Sergey Levine, Hugo Larochelle, and Yoshua Bengio. Recall traces: Backtracking models for efficient reinforcement learning. In International Conference on Learning Representations, 2019.
[7] Nathan Grinsztajn, Johan Ferret, Olivier Pietquin, Philippe Preux, and Matthieu Geist. There is no turning back: A self-supervised approach for reversibility-awee aware aware reinforcement learning. In A. Beygelzimer, Y. Dauphin, P. Liang, and Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[8] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. 2017.
[9] Todd Hester, Matej Vecerík, Olivier Pietquin, Mare Lanctot, Tom Schaul, Bilal Piot, Dan Horgan, John Quan, Andrew Sendonaris, Ian Oshand, Gabriel Dulac-Arnold, John P. Agapiou, Joel Z. Leibo, and Audrunas Gruslys. Deep gelaarnnng from demonstrations. In AAAI, 2018.
10

---

[10] Jonathan Ho and Stefano Ermon.
Generative adversarial imitation learning. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29, 2016.
[11] Ahmed Hussein, Mohamed Medhat Gaber, Eyad Elyan, and Chrisina Jayne. Imitation learningg survey of learning methods. ACM Comput. Surv., 50(2), apr 2017. ISSN 0360-0300.
[12] Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. In Advances in Neural Information Processing Systems, 2019.
[13] Beomjoon Kim, Amir-massoud Farahmand, Joelle Pineau, and Doina Precup. Learning from limited demonstrations. In C. J. C. Burges, L.. Bottou, M. Welling, Z. Ghahramani, and K. Q. Weinberger, editors, Advances in Neural Information Processing Systems, volume 26. Curran Associates, Inc., 2013.
[14] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learming for offline reinforcement learning. Advances Neural Information Processing Systems, 33: 1179-1191, 2020.
[15] Hang Lai, Jian Shen, Weinan Zhang, and Yong Yu. Bidirectional model-based policy opti- mization. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 5618-5627. PMLR, 13-18 Jul 2020.
[16] Michael Laskey, Jonathan Lee, Roy Fox, Anca D. Dragan, and Ken Goldberg. DART: noise injection for robust imitation learning. In 1st Annual Conference Robot Learning. CoRL 2017, Mountain View, California, USA, November 13-15, 2017, Proceedings, volume 78 of Proceedings of Machine Learning Research, pages 143-156. PMLR, 2017.
[17] Kimin Lee, Younggyo Seo, Seunghyun Lee, Honglak Lee, and Jinwoo Shin. Context-aware dynamics model for generalization in model-based reinforcement learning. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 5757-5766. PMLR, 13-18 Jul 2020.
[18] Yuping Luo, Huazhe Xu, and Tengyu Ma. Learning self-correctable policies and value func- tions from demonstrations with negative sampling. In International Conference on Learning Representations, 2020.
[19] Anirudha Majumdar and Russ Tedrake. Funnel libraries for real-time robust feedback motion planning. The International Journal of Robotics Research, 36(8):947 9822 2017.
[20] Matthew T. Mason. The mechanics of manipulation. In IEEE International Conference on Robotics and Automation, 1985.
[21] Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland. Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540):529-533 2015.
[22] Ashvin Nair, Bob McGrew, Marein Andrychowicz, Wojeiech Zaremba, and Pieter Abbeel. Over- coming exploration in reinforcement learning with demonstrations. In 2018 IEEE International Conference on Robotics and Automation (ICRA), pages 6292-6299. IEEE, 2018
[23] Suraj Nair, Mohammad Babaeizadeh, Chelsea Finn, Sergey Levine, and Vikash Kumar. Trass: Time reversal as self sulervuipn In 2020 IEEE International Conference on Robotics and Automation (ICRA), pages 115-121, 2020.
[24] Soroush Nasiriany, Vitchyr Pong, Steven Lin, and Sergey Levine. Planning with goal- conditioned policies. In NewrlPS, pages 14814-14825, 2019.
[25] Takayuki Osa, Joni Pajarinen, Gerhard Neumann, J. Andrew Bagnell, Pieter Abbeel, and Jan Peters. An algorithmic perspective on imitation learning. Foundations and Trends@ in Robotics, 7(1-2):1-179 2018. ISSN 1935-8253.

---

[26] Matthias Plappert, Marcin Andrychowicz, Alex Ray, Bob McGrew, Bowen Baker, Glenn Powell, Jonas Schneider, Josh Tobin, Maciek Chociej, Peter Welinder, Vikash Kumar, and Wojeiech 7aremba. Multi-goal reinforcement learning: Challenging robotics environments and request for research, 2018.
[27] Dean A. Pomerleau. Alvinn: An autonomous land vehicle in a neural network. In D. Touretzky, editor, Advances in Neural Information Processing Systems, volume 1. MlorannKauffnn 1989.
[28] Antonin Raffin, RI baselines3 zoo. https://github..cmmm/ www.///DDDDD/////aaaaaeeeeeeee 2020.
[29] Nasim Rahaman, Steffen Wolf, Anirudh Goyal, Roman Remme, and Yoshua Bengio. Learning the arrow of time for problems in reinforcement learning. In International Conference on Learning Representations, 2020.
[30] Aravind Rajeswaran, Vikash Kumar, Abhishek Gupta, Giulia Vezzani, John Schulman, Emanuel Todorov, and Sergey Levine. Learning Complex Dexterous Manipulation with Deep Reinforce- ment Learning and Demonstrations. In Proceedings of Robotics: Science and Systems (RSS), 2018.
[31] Siddharth Reddy, Anca D. Dragan, and Sergey Levine. Sqil: Imitation learning via reinforcement learning with sparse rewards. In International Conference on Learning Representations, 2020.
[32] Stephane Ross and Drew Bagnell. Efficient reductions for imitation learning. In Yee Whye Teh and Mike Titterington, editors, Proceedings of the Thirteenth International Conference Artificial Intelligence and Staristics, volume 9 of Proceedings of Machine Learning Research, pages 661-668, Chia Laguna Resort, Sardinia, Italy, 13-15 May 2010. PMLR.
[33] Stephane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Geoffrey Gordon, David Dunson, and Miroslav Dudík, editors, Proceedings of the Fourteenth International Conference Artificial Intelligence and Statistics, volume 15 of Proceedings of Machine Learning Research, pages 627-635, Fort Lauderdale, FL, USA, 11-13 Apr 2011. PMLR.
[34] Harsh Satija, Philip Amortila, and Joelle Pineau. Constrained Markov decision processes via backward value functions. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 8502-8511. 1. PMLR, 13-18 Jul 2020.
[35] Stefan Schaal et al. Learning from demonstration. Advances in neural information processing systems, pages 1040-1046, 1997.
[36] Julian Schritwiesser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Si- mon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604 609, 2020.
[37] Yannick Schroecker, Mel Vecerik, and Jon Scholz. Generative predecessor models for sample- efficient imitation learning. In International Conference on Learning Representations, 2019.
[38] David Silver, Aja Huang, Christopher J. Maddison, Arthur Guez, Laurent Sifre, George van den Driessche, Julian Schrittwieser, loannis Antonoglou, Veda Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy Lil- licrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel. and Demis Hassabis. Mastering the game of go with deep neural networks and tree search. Nature, 529:484-503 2016.
[39] Wen Sun, Arun Venkatraman, Geoffrey Gordon, Byron Boots, and J. Andrew Bagnell. Deeply AggreVaTedD VaTeD:: Differentiable imitation learning for sequential prediction. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 3309-3318, PMLR, 06-11 Aug 2017.
[40] O. Tange. Gnu parallel the command-ling power tool. :login: The USENIX Magazine, 36(1): 42-47, Feb 2011.
12

---

[41] Pierre Thodoroff, Audrey Durand, Joelle Pineau, and Doina Precup. Temporal regularization for markov decision process. In S. Bengio. H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
[42] Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: physics engine for model-based control. In 2012 IEEE/RSJ International Conference Intelligent Robots and Systems, pages 5026-5033. IEEE, 2012.
[43] Mel Vecerik, Todd Hester, Jonathan Scholz, Fumin Wang, Olivier Pietquin, Bilal Piot, Nicolas Heess, Thomas Rothörl, Thomas Lampe, and Martin Riedmiller. Leveraging demonstrations for deep reinforcement learning on robotics problems with sparse rewards. arXiv preprini a7Xiv:/707.08817, 2017.
[44] Arun Venkatraman, Martial Hebert, and Andrew Bagnell. Improving multi-step prediction of learned time series models. In Proceedings of the Twenty-Ninth AAAI Conference on Artificial Intelligence, AAAI' 15, page 3024-3033 3030. AAAI Press, 2015. ISBN 0262511290
[45] Jianhao Wang. Wenzhe Li, Haozhe Jiang. Guangxiang Zhu, Siyuan Li, and Chongjie Zhang. Offline reinforcement learning with reverse model-based imagination. Advances in Neural Information Processing Systems, 34:29420-29444 2021.
[46] Ziyu Wang. Josh Merel, Scott E Reed, Nando de Freitas, Gregory Wayne, and Nicolas Heess. Robust imitation of diverse behaviors. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[47] Keyulu Xu, Mozhi Zhang, Jingling Li, Simon S Du, Ken-ichi Kawarabayashi, and Stefanie Jegelka. How neural networks extrapolate: From feedforward to graph neural networks. arXiv preprint arXiv:2009. 1848, 2020.
[48] Tao Yu, Cuiling Lan, Wenjun Zeng, Mingxiao Feng, Zhizheng Zhang, and Zhibo Chen. Playvir- tual: Augmenting cycle-consistent virtual trajectories for reinforvement learning. In A. Beygelz- imer, Y. Dauphin, P. Liang and Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[49] Shangtong Zhang, Vivek Veeriah, and Shimon Whiteson. Learning retrospective knowledge with reverse reinforcement learning. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 19976-19987. Curran Associates, Inc., 2020.
[50] Yuke Zhu, Ziyu Wang, Josh Merel, Andrei Rusu, Tom Erez, Serkan Cabi, Saran Tunyasuvu- nakool, János Kramár, Raia Hadsell, Nando de Freitas, and Nicolas Heess. Reinforcement and imitation learning for diverse visuomotor skills. In Proceedings of Robotics: Science and Systems, Pittsburgh, Pennsylvania, June 2018.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [No]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you experiments...
13

---

(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URL)? [Yes] Descriptions of the implementation are provided in the main text and supplementary material and the code and data are released as a public repository.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] Briefly in Section 4, and in full in the supplementary material. (c) Did you report error bars (e.g.. with respect the random seed after running experi- ments multiple times)? [Yes] See Section 5.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] See supplementary material.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [No]
(c) Did you include any new assets either the supplemental material or as URL? [Yes]
See Section 4.2 for the code repository URL.
(d) Did you discuss whether and how consent obtained from people whose data you're
using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "Zxjhh_SBsQo", "reviewer": "Reviewer_rhnT", "summary": "This work introduces a novel method that incorporates a backward model into the offline imitation learning process to improve the model robustness against covariate shifts. The generated demonstrations from the backward model would train the policy to recover from the error states and back to the main distribution that it is familiar with. The motivation is clear and the idea of using offline data augmentation is a promising direction for improving the robustness of imitation learning algorithms.", "questions": "1. Why train the backward model together with the policy model? As an offline method, I thought the easiest way is to first train the backward model alone and then train the policy alone. I hope the authors could give more details and motivation about the training pipeline.\n2. What’s the observation used in the Fetch experiment? I assume it is the proprioception / joint angles of the Fetch arm. Do you give the target pose of the object? Is the proposed method limited to low-dimensional observations? Any discussion about potential high-dim observation inputs like images would be great.\n3. How does the method perform if the given demonstration is sub-optimal? Let’s say the demonstration is collected by a human expert and there are imperfect motions. Will it be a challenge for this method?", "limitations": "1. Bi-directional state transition assumption.\n2. Low-dimensional observation space.\n3. Might be hard to handle sub-optimal demonstrations.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n1. The paper is well-written. The problem setting and technical terms are clearly introduced.\n2. The idea of learning a backward model to generate more offline demonstrations is novel. Most prior imitation learning works only focus on the method part, I believe a good way of augmenting the training data is also an important direction to improve the robustness of offline imitation learning.\n3. The idea of decomposing the backward model into two functions is a good way to model the different features of the action generation and state transitions. To handle the multi-model distribution of action generation, a conditional Gaussian probabilistic model is used, which could further generate diverse trajectories by sampling the noise. I found these architectural designs promising and reasonable.\n4. The experiment results in both robot and maze environments show the promising advantages of the proposed approach.\n\nConcerns:\n1. One assumption is missing: this method assumes that the state transition in the environment is bidirectional. While in many real-world scenarios, the forwarding trajectory is not reversible which might make the generated traces impracticable in the actual environment. The proposed method might struggle when the environment has many bottleneck state transitions that are not reversible. But I do agree this is out of the main focus of this work.\n2. It seems like the experiment tasks are more like some proof of concept. I'm a little bit concerned about the performance in high-dimensional and complex tasks. For example, what if the observation is an image?\n3. I do agree generating more trajectories from given demonstrations could improve the robustness of the learned policy. But it is still unclear to me why the model could generate novel initial states and train the policy to recover from it. In the Fetch experiments, it seems like  BMIL solved many cases with unseen initial states. But these states shouldn’t be generated by the backward model since it is unseen in the given demonstrations. Then, why the policy model could handle it so well?", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "RqWeY14Ikpg", "reviewer": "Reviewer_VFop", "summary": "This paper seeks to improve the robustness of traditional behavior cloning, by aggregating backward traces from the demonstration to nearby states into the dataset. They do so by learning a factored backwards dynamics model, and rolling out this dynamics model from demonstration states to retrieve backward traces. By training a subsequent learner policy on original demonstrations and backward traces that lead the agent back to the demonstrated states, they improve the robustness of the learned policy on 4 different domains. \n", "questions": "In addition to the primary in-line questions mentioned in the above section, here are a few other questions I have - \n1. In Line 148-149, the paper mentions irrecoverable states. The reviewer agrees that the assumption of existence of recovery actions is reasonable outside of such irrecoverable states. However, does this preclude the use of the approach in domains where irrecoverable states exist? For example, could a backwards model be queried at states nearby irrecoverable ones, and be made to be more robust near such states? \n2. How do the authors balance the backwards model B being accurate w.r.t the forward policy, with its ability to generate diverse new traces? (as mentioned in lines 163-164). \n3. How reasonable are the assumptions of Lines 179-180, of “”If the reverse time model B is accurate and the previous action generator BA gives sufficiently diverse actions”? \n4. In Line 195, is the definition of $a’$ meant to be $a’ \\sim \\mathcal{N} ( \\mu, \\sigma), \\mu = [ a - k, a + k]$, rather than the specified equation? \n5. In line 191, why is it necessary for the action-selection strategy to cover the expert actions? Is this to maximize robustness? What happens to robustness if the action selection distribution and the expert actions only intersect? \n6. In Line 198, the paper states “As we do not wish to deviate too much from the demonstration data, we perturb only the first action when $s’ \\in \\mathcal{D}_{demo}$”. Why is this desirable? Why don’t we want to stray far away from the demonstration data? \n7. In Lines 203-205, the paper states “In addition to using relatively short rollouts that limits compounding errors, we hypothesize that this is these traces do not necessarily need to be accurate in order to benefit the policy. It may be that having the general correct direction contained in the traces is sufficient for the policy to eventually reach states in the demonstration data”. By this logic, isn’t it sufficient to have an arbitrary action selection strategy, rather than one that relies on the backwards dynamics? \n8. How are the forward dynamics models trained in Line 298 used? \n", "limitations": "For limitations, please view the section on weaknesses above.\n", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n\n** Clarity: **\n\n1. Overall, the paper is written very well. The problem tackled is well motivated, the approach is described in detail and most design choices are well justified. The contribution is well contextualized in the context of prior work, and the experimental results are also described in sufficient detail. Two particular ways in which the paper is written that should be appreciated are - \n\n    a. The paper does not exaggerate their claims beyond what their experimental results show. \n\n    b. The work also pre-emptively answers several questions that naturally arise during the discussion of their work (such as the accuracy of the backward dynamics model, and the effect it has on robustness of the learnt policy).\n\n** Quality: **\n\n2. Overall, the paper is also technically sound. \n\n    a. The paper honestly describes their claims (more robustly solving a set of tasks that current IL algorithms can currently solve), and their claims are supported by the experimental results they present. \n\n    b. The choices of methods are well justified, and for the most part, the paper provides intuition about the functioning of their method (such as the dynamics of accuracy of the backwards model in training, and its relation to how robustly the policy learns). \n\n    c. I also appreciate that the paper is very upfront about which specific problem setting they tackle, and when their approach would, and would not, be useful.\n\n### Neutral\n\n** Originality **\n1. The originality of the paper is fair. The paper adopts a very similar approach to prior work [41], but applied to the imitation learning domain rather than the reinforcement learning domain as in [41]. This suggests the approach itself is not entirely novel, and only differs in some of the lower-level implementation details. This; however, is acceptable - the paper is upfront about this similarity to prior work, but clearly discusses why solving a similar problem in the imitation learning domain is important, why current methods that do so are limited (such as DAgger being expensive to run, etc.). As a result, one must give credit to the paper for identifying the applicability of this approach to the imitation learning setting. As a result, the overall originality of the paper is fair. \n\n### Weaknesses\n\n** Significance **\n1. While the experimental results presented are technically sound, they are presented in somewhat limited settings. \n\n    a. The tasks that the paper evaluates their approach on are quite simplistic tasks, such as solving a maze or controlling the Fetch end-effector to reach / push / place objects. In these tasks, “robustness” of the policy is a nice property, but not crucial to success or necessary for safety. The paper could benefit from evaluating their approach against baseline approaches on a more diverse set of tasks that are inherently harder to solve. Examples of such tasks could be more complex manipulation tasks, tasks in the joint space, continuous navigation tasks, and more interestingly, tasks where dynamics matter as well, such as continuous control of cars or legged agents. While I appreciate the authors mention this is a plan in future work, the impact of the paper in its current form is a tad underwhelming, and could be greatly improved by such evaluation on diverse tasks. \n\n2. Applicability of approach. The paper could also benefit from both discussing, and experimenting on, a broader settings (in addition to other tasks as mentioned above). The paper currently states such settings are not considered, but it is unclear why. Some such settings include - \n\n    a. For some reason, the paper only explores robustness with respect to spatial dimensions (such as joint coordinates, and not velocities, as mentioned in Lines 143-144). It is unclear why they restrict their discussion to this setting. Being able to recover to arbitrary states, and not just spatial dimensions, would be a big improvement to the paper. Also, several interesting questions then arise - does one need to follow a different action selection strategy for such dimensions? Does one need to train the backwards dynamics model any differently for such cases? Are backwards dynamics models as well conditioned when joint velocities or other non-linear dynamics are concerned? \n\n    b. Another interesting setting is when there are additional demonstrations available. Quantifying how much the robustness improves from the proposed approach with respect to number of additional demonstrations / some measure of coverage of state / action space of the demonstrations would be interesting to see, and useful in assessing the impact of the paper.\n\n    c. Another interesting area (but one that is likely better suited for future work) is when there is also access to computationally expensive interactive experts and demonstrations. Could the proposed approach be combined with interactive feedback to verify or improve robustness of the learner?\n\n### Note to authors\nI strongly urge the authors to consider evaluating their approach on the above suggested domains / settings, as I believe it would contribute to making their paper much stronger. Currently, my biggest reservation with the paper is these evaluations, but should these be improved, I think this could make for a strong paper. \n", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Pi8XgCKORgU", "reviewer": "Reviewer_L7Ps", "summary": "The paper presents a simple method to robustify policies that are trained from a small number of demonstrations using supervised learning (i.e., behavior cloning). The key idea is to synthesize recovery trajectories that are not included in the training set. This is done by generating possible short-horizon trajectories that end at the demonstrated state using a learned backward dynamics model $p(s_t, a_t| s_{t+1})$. The trajectories are then used as additional data to train the behavior cloning policy. The backward dynamics models are encouraged to generate diverse trajectories in order to enrich the set of out-of-distribution states that the policy can handle. Both the policy and the backward dynamics model are trained from a small number of demonstrations. The method is shown to achieve better robustness compared to naive behavior cloning method and a competitive offline IL baseline (VINS) on two simulated environments. ", "questions": "Please refer to the main comments listed above.", "limitations": "The paper does not mention any limitation of the proposed method. Like mentioned in the comments above, I hope the authors could stress test the method and discuss the limitations of the method both conceptually and empirically.\n", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall I enjoyed reading the paper. The method is conceptually simple and easy to apply on top of existing offline imitation learning methods. The empirical results also seem to support the main claims of the paper. \n\nMy main concern about the paper is the scope and the details of the experimental evaluation. \n\n- The evaluation tasks are rather simple both in terms of the nature of the tasks (reaching / pushing & 2D nav) and the types of dynamics involved. Hence it is unclear to me whether the method is still effective for more complex tasks such as Franka Kitchen and dexterous manipulation (Fu et al., 2020). Given the ease of implementation and the fact that these environments have expert data available, I’d strongly encourage the authors to stress test the proposed method on these harder tasks. \n- How is the forward dynamics model baseline implemented? Specifically, how does the policy make use of the rollout generated by the forward dynamics model?\n- How does the accuracy of the backward dynamics model influence the final robustness performance? I’d suggest a simple ablation study that injects different levels of noise to the learned backward dynamics model. \n- According to Table C.1, even BC can achieve near perfect performance in both Fetch and Maze tasks with only 5-20 demonstrations. Unless I’m missing something, this seem to directly contradict with the results in prior works (e.g., in D4RL) that require a few orders of magnitude more data to achieve reasonable performance. I hope the authors can explain this discrepancy or point out my misunderstanding. \n", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept