KitchenShift: Evaluating Zero-Shot Generalization of Imitation-Based Policy Learning Under Domain Shifts
Abstract
Humans are remarkably capable of zero-shot generalizing while performing tasks in new settings, even when the task is learned entirely from observing others.
In this work, we show that current imitation-based policy learning methods do not share this capability, lacking robustness to minor shifts in the training environment.
To demonstrate these limitations of current methods, we propose a testing protocol that new methods may use as a benchmark.
We implement and evaluate
KitchenShift, an instance of our testing protocol that applies domain shifts to a realistic kitchen environment.
We train policies from RGB image observations using a set of demonstrations for a multi-stage robotic manipulation task in the kitchen environment.
Using KitchenShift, we evaluate imitation and representation learning methods used in current policy learning approaches and ﬁnd that they are not robust to visual changes in the scene (e.g., lighting, camera view) or changes in the environment state (e.g., orientation of an object).
With our benchmark, we hope to encourage the development of algorithms that can generalize under such domain shifts and overcome the challenges preventing robots from completing tasks in diverse everyday settings.
Introduction
The real world is ever-changing, with each moment bringing new, visually rich scenes for an agent to see and interact with.
Consider a person going to a neighbor’s kitchen for the ﬁrst time to help cook dinner.
Humans have a profound ability to maneuver such never-before-seen environments, transition effortlessly between similar domains, and generalize within narrow task speciﬁcations.
These abilities lie in stark contrast to the current abilities of robots and policy learning algorithms.
How can we narrow this gap?
We argue that the development of better algorithms that could be deployed in real-world robotics settings has been bottlenecked by the simulation settings and evaluation schemes that are currently used.
Established benchmarks, such as video game tasks from Atari [9] and simple control tasks from OpenAI Gym [10] or DeepMind Control Suite [93], evaluate policies under the same environment settings in which they are trained.
Our evaluation scheme is designed to test robustness to visual and environmental changes using realistic out-of-distribution domain shifts.
Speciﬁcally, we provide a robotic agent with a set of expert demonstrations and train it to perform a multi-stage manipulation task in a realistic kitchen environment.
The policy is then evaluated in an unseen scene, which is created by applying one type of domain shift to the training environment.
This situation is designed to mirror a scenario that humans excel at but current algorithms do not.
In this work, we make the following contributions:
[Figure] • We deﬁne an evaluation protocol to determine the robustness of policies to minor, out-ofdistribution domain shifts to the environment, which a human could adapt to.
• To use this evaluation protocol, we introduce KitchenShift, which adds functionality for applying 7 types of domain shifts to the realistic kitchen simulation from Gupta et al.
• We demonstrate KitchenShift by benchmarking different imitation-based policy learning and representation learning methods, and ﬁnd that simple behavioral cloning outperforms other approaches evaluated at zero-shot generalizing under domain shifts.
Overall, we observe that current approaches are not robust to realistic out-of-distribution domain shifts.
In order to eventually deploy robotic agents into everyday life, we will need approaches that are robust and generalize to out-of-distribution settings.
We release KitchenShift to encourage the development of future algorithms that can overcome such domain shifts.
Evaluating Policies under Domain Shifts with KitchenShift
In this section, we describe KitchenShift, the implementation of our evaluation paradigm.
We divide the section into high-level protocol (Section 2.1), domain shifts (Section 2.2), and policy evaluation (Section 2.3).
In Appendix B, we discuss the choice of the kitchen environment we use.
2.1 Evaluation protocol
“Domain” and “distribution” shifts are frequently used interchangeably.
In this paper, we use “domain shift” to refer to a change in the environment, world, or domain D .
We do not take “domain shift” to include changes to the distribution or set of tasks that the policy learns.
For our evaluation protocol, we train the policy on a task T train in a new training domain D train , with access to a limited set of demonstrations.
These demonstrations reduce the need for exploration, improving sample efﬁciency and lowering compute requirements.
We then evaluate how the policy performs the training task T train in an unseen set of domains {D shifted } by modifying the evaluation environment.
If the policy has learned general behavior, then it should be able to adapt to the new testing domain and perform the task under such domain shifts.
We describe and provide a set of domain shifts that can be applied individually to independent parts of the environment, which an average human would be able to overcome and still perform the training task successfully.
Though we focus on a single environment in this work, we note that this setup can also be used to evaluate policy adaptation and transfer to new domains.
For instance, a policy could be pretrained on a set of tasks in RoboSuite [116] then ﬁnetuned and tested on KitchenShift, our implementation of the evaluation protocol we propose, which we proceed to describe.
2.2 Domain shifts in KitchenShift
Building off the kitchen environment released by Gupta et al.
[33], we make changes to the simulation settings (see Section E) and add functionality to modify the environment, which we refer to as
KitchenShift.
The domain shifts that we evaluate are not only limited to visual changes such as background context [36, 88], colors, lighting, textures, or camera pose [41, 113, 37, 25].
In Figure (Right), we visualize the types of domain shifts available in KitchenShift that can be applied to create different testing scenes.
We categorize seven types of domain shifts, based on changing (a) object instance; (b) object layout; (c) texture; (d) camera view; (e) lighting; (f) robot state; (g) object state.
For (a) object instance, we change the microwave and kettle assets 2 .
For (b) object layout, we change the position and orientation of different objects.
For (f) robot state, we change the initial joint positions of the robot manipulator.
For (g) object state, one object in the scene is changed to its ﬁnal goal state.
The training scene, which is also the domain used by the demonstrations, is shown in
Figure 1 (Left).
Our goal is to isolate and analyze differences in the performance of existing algorithms, to inform where these methods may be improved.
To create testing environments that are explicitly outside of the training distribution, we apply a single type of domain shift to the training environment.
If a human was asked to perform the task in the training domain, they would subsequently be able to adapt and succeed in such shifted environments as well.
Humans are able to learn from a narrow task speciﬁcation and training domain, while easily adapting to shifts in the scene.
We do not apply domain randomization by training the policy on the set of domain shifts encountered during evaluation.
Our purpose is to evaluate policy robustness when encountering unseen, out-of-distribution scenes.
2.3 Policy evaluation
The training task that the policy is asked to learn is multi-stage and involves interacting with a set of at most four objects in the environment.
In Appendix C, Figure 4, we visualize frames selected from an expert demonstration for the training task that involves manipulating the [ microwave,kettle,switch,slide ] objects in the scene to the ﬁnal goal state.
We evaluate policy performance in the training domain and average performance across testing domains within each category of domain shift.
We also average across categories of domain shifts to report the overall performance across all testing domains, as a measure of the robustness and generalization of the learned policy.
We exclude the (g) object state category from this measure because the object changed may be involved in the task and would improve the average step completion reported, inﬂating policy performance.
For instance, one of these domain shifts opens the microwave door to the desired goal state (see Figure 1 (Left)), allowing a robust policy to skip interacting with the microwave.
The (g) object state category is still evaluated separately in Section 3, Figure 2.
Experimental Results
KitchenShift is available at github.com/etaoxing/kitchen-shift .
We provide details on the methods we evaluated and experimental details in Appendix C and Appendix D, respectively.
In
Appendix E, we discuss and validate the random perturbations we apply when initializing the simulation environment, in contrast to the original environment code.
We compare representation learning and imitation learning methods based on behavioral cloning, in the training domain and in unseen testing environments with speciﬁc types of domain shifts applied.
Performance of policies trained with different methods is reported in Table 1, and results show that learning-based methods are less successful when evaluated under domain shifts.
Evaluating policies with KitchenShift underscores the pitfalls of learning-based policies when testing generalization to out-of-distribution settings, including even minor deviations from the training domain.
To further interpret these results, we show policy performance separated by each category of domain shift in Figure 2.
Intuitively, Demo playback is not affected at all by (c) texture, (d) camera view, or (e) lighting, as these are visual-only changes.
On the other hand, Demo playback cannot account for changes to the initial robot state of the trajectory or environment, while learning-based methods show some robustness to this.
We ﬁnd that BC (MSE) and BC ( β -VAE) drastically overﬁt to the training distribution and fail to generalize out-of-distribution under minor domain shifts.
In comparison, BC which parameterizes a policy as a logistic mixture distribution, generalizes better to unseen scenarios.
We also ﬁnd that current representation learning approaches are not helpful in this imitation-based policy learning setting, which corresponds with recent work by Chen et al.
Note that the (g) object state category cannot be directly compared to the other types of domain shifts, as it sets one object in the scene, that may be involved in the task, to the desired goal state.
The
[Table]
" # $ % $ % & ’ # ( ) * + , - . ! )
$ % / !
+ ( ) * + , - . ! )
# 5 - " # ) 6 $ -’ - ( ) 0$ & 8 !
$ % & ’ 9 ( ) " * + * ! ) / !
# ! - ’ & ( ) * + , - . ! ) / !
# ! - : :;<
= =;< > >;< ? ?
;< @
AB
AB)’CDE(
AB)7F)$%6;)41%#5$./
AB)’ !
GHIE(
AB)’D$5D$#5J)9200)#2&;())
AB)’IKBJ)"#%4*5)/8$9!
AB)’L3-4)"#%4*5)-5+-44$%&())))
M-5*)N0#1+#.O
[Figure] [Figure] [Figure] [Figure]
[Figure]
categories that shift (a) object instance; (b) object layout; (f) robot state; and (g) object state, which are not just visual changes but modify the underlying scene, pose the most interesting challenges.
One approach to overcome these challenges for future work may be to replace implicitly learned representations with learned object-centric representations [21, 57].  Related Work
We cover prior work on benchmarking policy learning in Appendix A.
Addressing domain shift Current approaches for dealing with domain shifts, such as the gap between simulation and the real world, involve system identiﬁcation to tune simulator parameters [22] or by applying large amounts of domain randomization [95, 2, 5]. While such approaches may work for constrained lab settings or in the data-rich area of computer vision, they require large engineering effort to deploy for robot learning. Collecting large amounts of real-world data with robots is also time and capital intensive [56, 19]. To deploy robots that can learn and adapt directly in the real world, we need drastic improvements in the robustness of learning approaches to out-of-distribution changes.
Distribution shift [90, 48, 32, 101] has been studied before in the context of supervised learning such as image classiﬁcation, but domain shift for policy learning is challenging and under-explored. To the best of our knowledge, single domain generalization [98, 114, 74, 100] has not been studied before in the context of policy learning. Several works have evaluated visual representation learning [52, 36, 89] and data augmentation [53, 75, 109, 108, 37] alongside policy learning to improve generalization, however these methods train with visual inputs on simple control tasks [93] or Procgen [17], a video game environment. In this work, we show that policies trained with these representation learning approaches still suffer from domain shifts in a more challenging and realistic kitchen environment.
Discussion
We hope KitchenShift, with the evaluation protocols we developed and tested, will cast light on the fragility of imitation-based policy learning approaches.
Our results show that current methods are not robust under minor visual and environmental domain shifts to which humans could effectively adapt.
We also evaluate several representation learning approaches and ﬁnd that these methods are still unable to overcome the challenge of generalizing with realistic visual scenes in the limited data domain.
To move toward robots that can learn directly in the real world and adapt to new, unseen situations, the robotics community will need methods that tackle out-of-distribution domain shifts head-on.
Future directions Future work may evaluate transfer and generalization of policies by pretraining with different simulators, followed by training and evaluation using KitchenShift.
This is an exciting direction, as transfer learning has had signiﬁcant success in computer vision, but how to transfer policies between environments or different robot morphologies remains an open question.
On a related note, test-time training [91, 7], which adapts models online rather than ﬁxing parameters for evaluation, presents an interesting path forward.
Ideally, learning-based methods deployed in the real world are able to continuously learn and adapt to new settings.
Relaxing some of the available data assumptions, such as imitation directly from a demonstration video without access to the trajectory actions, also poses an interesting line of work.
This could involve imitation learning by watching YouTube videos of humans, for which there will be more challenging domain shifts between the real-world videos and the kitchen simulation.
Limitations We evaluate representation learning methods primarily alongside imitation learning, due to time and compute constraints.
Such methods may perform differently when trained online with reinforcement learning.
Still, we do not expect changes to the ﬁnal results, as these methods should have an easier time learning in the simpler, stationary imitation setup.
We also only test our evaluation protocol using one kitchen environment, but we expect our results would transfer to other simulations, as the underlying policy learning methods would still suffer from similar challenges.
Broader impact In this paper, we evaluate policy robustness to minor, out-of-distribution domain shifts to which a human could easily adapt.
Understanding these limitations will aid in the development of new algorithms that can deploy on robots in the real world.
Understanding where methods fail is also important in determining the expected behavior of policies for robot safety [3, 30].
On the other hand, improving policy learning algorithms towards human-like generalization capabilities may lead to large-scale automation, disrupting the workforce across many sectors [12].
For the moment, our work is far from these impacts, and we hope that the positive implications of this line of research will come to fruition in parallel with regulation and caution, to mitigate negative consequences.
A Extended Related Work
Benchmarking policy learning Random search algorithms have been shown to solve [62] the established benchmarks such as low-dimensional control tasks from the OpenAI Gym [10] or
DeepMind Control Suite [93, 94].
Furthermore, RL algorithms which have frequently been developed and evaluated on these tasks are known to be brittle, sample-inefﬁcient, and difﬁcult to apply in the real world.
Video game-like environments [46, 97, 16, 34, 80] such as Atari [9] have also been used to evaluate policy learning, but these environments do not have realistic graphical assets and do not simulate real-world physics.
While procedurally-generated game environments [44, 51, 17, 6] have been proposed to evaluate policies in unseen testing environments, the test environments are usually in-distribution to the train settings and difﬁcult to apply realistic instances of domain shifts in.
Many environments have also been introduced that involve robots for object manipulation using a manipulator arm [71, 33, 111, 42, 116, 59, 45, 67, 55] or dexterous hand [76, 5], as well as mobile agents in more complex settings such as home environments [11, 105, 81, 73, 103, 29, 49, 87, 23, 86, 104, 28, 82, 92].
We note that planar manipulation tasks of rigid [112] and even deformable [83] objects, where the robot is given a top-down view of a tabletop, can be nearly or completely solved by methods that can take advantage of such spatial structure.
As such, we avoid selection of these tasks for this benchmark.
Additionally, the latter types of home simulation environments pose an interesting and challenging set of tasks for policy learning, compared to more basic scenes that only contain a robot and objects on a solid-colored surface.
However, these simulators are difﬁcult to set up, with many layers of software complexity.
State-of-the-art RL algorithms are also very far from learning on such long-horizon manipulation and navigation tasks.
In the search for better evaluations, there has been a consistent stream of benchmarks [20] released recently, and in particular environments for the evaluation of policy learning.
In this work, we do not introduce a new environment, but instead add new functionality to an existing environment [33] that has already been used by the community [27, 50, 69, 70, 58].
The task we train policies on can be viewed as a kind of rearrangement task [8], where a ﬁxed robot arm is asked to manipulate the objects in the scene to some given goal state.
We do not use a mobile manipulator and circumvent the challenge of combining navigation and manipulation in this work.
Commonly, RL and robotics methods either report performance on the training domain or assume that the train and test sets are drawn from the same distributions, following standard assumptions and practices with training neural networks.
We take the stance that deploying robotic agents in the real world requires policy learning approaches that can generalize or adapt to out-of-distribution testing scenes, which we evaluate in simulation with KitchenShift in this work.
For causal learning in robotics, CausalWorld [1] introduces task generators for parameter variation of cuboid block manipulation with a three ﬁnger robot, though their setup could be modiﬁed to evaluate zero-shot generalization.
Our work is similar to ManiSkill [63] and MAGICAL [96], which leverage demonstrations and evaluate policy generalization.
However, MAGICAL uses a visually simplistic 2D world, and ManiSkill is focused on robust object-level manipulation in an empty background scene.
In contrast, we evaluate out-of-distribution policy generalization across 7 categories of domain shifts in a realistic kitchen scene.
We ﬁnd that BC alone performs better than other methods, which corresponds with results from other recent work [63, 43, 61, 102, 26] that also evaluate on more challenging tasks.
B Choice of Environment
To evaluate policy learning, we chose the kitchen environment, task deﬁnitions, and demonstration trajectories from Gupta et al.
This environment is visually realistic with challenging multi-stage tasks involving a robot manipulator.
Demonstrations given by human VR teleoperation of the robot are also available, compared to environments which use carefully designed policies to generate demonstrations or even ones without any demonstration data.
Furthermore, this environment has already been used by the community [27, 50, 69, 70, 58].
MuJoCo has been acquired and open-sourced by DeepMind, which makes KitchenShift accessible for the community.
There is also active development on other simulator backends such as NVIDIA Isaac Sim, [60], and in the near future, KitchenShift could be ported over to use other simulations.
Prior work [95, 2, 5, 22] has used the MuJoCo simulator for the purposes of domain randomization, and we use similar functionality to apply domain shifts to the environment.
We also note that our domain shift evaluation protocol is not uniquely applicable to this particular kitchen environment.
Our evaluation paradigm could be applied to other environments such as recent home environment simulations [49, 86, 104, 92, 28], which also have a wide range of objects, texture assets, and scenes; though that is outside the scope of this work.
C Experimental Setup
Using domain shifts and the KitchenShift benchmark as described in Section 2, we evaluate the generalization of different imitation-based policy learning approaches.
As in Gupta et al.
[33], we report the average step completion, out of a maximum of four total subgoals.
Agents are given a small set of 19 expert demonstration trajectories of the training task.
All results reported are averaged across 4 random seeds.
Additional details can be found in Appendix D.
We prioritized methods with publicly available code that were easier to implement and tune, as well as ones that could take advantage of learning from demonstrations for improved sample efﬁciency.
The implementations we developed are open-sourced in a shared codebase 3 .
The repository also includes the full set of experiment logs and models used to generate the results in this paper.
We evaluate combinations of the following imitation learning and representation learning baselines.
C.1 Imitation learning
We use an implementation of behavioral cloning (BC) open-sourced by Dasari and Gupta [18], which parameterizes the policy as a discretized logistic mixture distribution [79, 59].
We include a comparison to standard BC trained with mean-squared error (MSE), which corresponds to maximizing the expected log likelihood under a Gaussian distribution [68].
We also experiment with jointly optimizing an auxiliary dynamics modeling objective to regularize policy learning, which was studied by Dasari and Gupta [18] and in other work as well [107, 54, 99, 106].
To support training with multi-task demonstration data, we also evaluate goal-conditioned BC (GCBC) [59], which conditions the policy on a goal image specifying the task to perform.
We use the goal relabeling strategy originally proposed by Hindsight Experience Replay (HER) [4], but we do not call a reward function to relabel rewards as in HER.
C.2 Representation learning
Reconstruction representations have been used extensively for perception in policy learning.
We experiment with different autoencoders [47, 77]: β -VAE [40], RAE [31], and σ -VAE [78].
For robotics applications, β -VAE is commonly used [65, 64, 72, 115].
Rybkin et al.
[78] recently proposed σ -VAE as an extension to β -VAE that automatically sets the β hyperparameter.
Furthermore, Yarats et al.
[110] ﬁnd that a model-free RL policy using a deterministic autoencoder (RAE) outperforms more complex model-based planning approaches that employed β -VAE, when the environment background contained noise distractors.
Building off the recent success of instance discrimination for self-supervised learning in computer vision [13, 38, 15], contrastive representations have been studied for policy learning.
We evaluate
ATC [89], a version of noise contrastive estimation [35, 66, 13, 38] that associates images within a temporal window and applies the random shift augmentation proposed in Yarats et al.
We also experiment with SimSiam [15], which associates two augmentations of the same image without using negative pairs, and use the larger set of augmentations from computer vision tasks.
The same four-layer ConvNet from Sekar et al.
[84] is used for the encoder network to compare different methods.
As a baseline, we ﬁx the initial random weights of the encoder network layers [89, 85], to validate that the policy leverages image observations to aid control.
We also experiment with a deeper encoder network architecture: the 15-layer network used in IMPALA [24].
C.3 Random actions and Demo playback
We also report performance on two basic, non-learning baselines.
For Random actions, we sample an action uniformly from the action space at each timestep.
For Demo playback, we randomly sample a demonstration and execute the actions given by the demonstration trajectory open-loop.
D Experimental Details
We render (256, 256) RGB image observations of the environment and resize them to (224, 224) as input to the agent, following conventions from computer vision with ConvNets for image classiﬁcation.
Image observations are embedded into a 128-dim latent vector by an encoder network.
We also provide the policy with a noisy observation of the robot’s proprioceptive state (36 dimensions): joint positions, joint velocities, as well as end effector position, orientation, and contact forces.
Policies are parameterized by a three-layer feedforward neural network with 256 units per layer.
The continuous action space is 9-dim, controlling the robot gripper and the arm’s joint velocities.
For methods that condition on goal context, we provide an image that speciﬁes the ﬁnal environment state (GCBC) and use the same encoder network as for the image observation.
All methods are trained with the Adam optimizer ( α = 0 .
001 , (cid:15) = 0 .
01 ) and a batch size of 128 episode steps for a total of 1e5 optimization iterations.
In early experiments, we trained policies for longer but found 1e5 iterations to be sufﬁcient to evaluate policies.
Beyond the single training task evaluation where the policy must interact with the [microwave,kettle,switch,slide] objects, we widen the training distribution to include demonstrations of other tasks.
All demos selected involve at least interacting with the microwave object, with 235 demos across a total of 12 multi-stage tasks.
Agents are trained for twice the number of optimization steps used in the single training task setup.
We train GCBC with this larger demonstrations dataset to compare to GCBC trained only with trajectories of the single task.
We used a university cluster to generate results; we ran four seeds at a time on 4-GPU (either Titan
X (Pascal) or 1080 Ti) machines with 32-vCPUs each.
Conﬁg ﬁles are available in our repository, containing hyperparameters used for all methods and experiments we ran.
E Impact of simulator initialization
The original kitchen environment code [33] resets and initializes the state of the underlying physics simulator in a deterministic fashion.
We found that this determinism inﬂates the performance of methods, as the policy is evaluated in the exact same simulation conditions, even with uniform noise added to proprioceptive observations given to the policy as per the original settings.
To illustrate this gap, we generate demonstrations in environments with deterministic initialization and with noise perturbations applied during initialization.
For the latter, we combine two types of perturbation: (i) small uniform noise added to the initial state when resetting the simulator; (ii) stepping the simulation a random number of times when resetting.
Perturbation (ii) further changes the simulator state because we do not compensate for gravity in the robot controller, so the unactuated robot arm can drift.
This noise perturbation only affects the initial state of the environment.
In Table 2, we show results of deterministic versus stochastic environment initialization when evaluating policies.
We report results with BC (state-vector) as a baseline which achieves similar performance as the original behavioral cloning results from Gupta et al.
[33] in environments with deterministic initialization.
Notably, evaluating in environments with the stochastic initialization reduces policy performance for all methods, most notably by 51% for BC ( β -VAE).
[Table] To ensure that we report on a reasonable metric for realistic conditions (the importance of which is highlighted by Henderson et al. [Table]
for all other experiments evaluate policies in environments with stochastic initialization.
Our purpose here is to show that a small amount of perturbation to the initial state signiﬁcantly reduces the reported performance of the policy.
F Domain shifts evaluated in KitchenShift