Action Inference by Maximising Evidence: Zero-Shot Imitation from Observation with World Models

Summary

This paper considers the setting where we have access to a dataset containing states and actions for pre-training and must then learn to solve a task given a new observation-only dataset from a downstream task. The key assumption here is that we have access to this action labelled dataset, and once we have it, we are better off training a world model than simply learning to label the new observation only trajectories with action labels (a la BCO and subsequently VPT). The proposed method is one of the first (as far as I know) to use forward dynamics models to implicitly model the action distribution and provides a nice alternative to IDM approaches for what is becoming an increasingly relevant problem setting. Given the relevance of the topic area, and the fact this exact thing has not been done before, I am voting for the paper to be accepted. My score would be increased if some additional experiments could be conducted, since these particular ones are relatively similar and also low impact in terms of their ambition.

Strengths

The strengths of this work are clear, it is a very relevant problem setting and this method is distinct vs. previous methods like BCO which rely on an IDM. In reality, this paper is essentially "Implicit action learning with forward dynamics models", and that has not been done before as far as I am aware. The method itself is fairly clearly presented, and the experiments are relatively clear with sufficient ablation studies. Finally, it is great to see limitations adequately discussed in the main body, which is surprisingly rare.

Weaknesses

Note that I have voted to accept, the following comments are not red flags but would likely improve the paper, and maybe make it possible to increase to a higher score. 1. The experiments are fairly mundane, and while scientific best practices appear to have been followed, there is a huge gap between what the authors claim to be working towards ("a single foundational world model") and what is actually shown (two DMC environments). It would be fantastic to see an example of this method at larger scale, even if the results are not state-of-the-art and there is only a single seed, for compute reasons. For example, this could be done using the dataset from VPT. The Minecraft images could be resized to make them smaller, and then it would be possible to use the DreamerV3 codebase (which was tested on MineCraft and runs on a single GPU) then see if it is possible to learn from the unlabelled MineCraft videos. If this works, it would drastically increase the impact of the paper, beyond being something mildly interesting for people who care about this specific topic, to something that catches people's eye across the field. 2. It seems slightly fishy to use a world model generated dataset to compare a world model based method and an IDM approach. It is possible there is some bias in the Dreamer or P2E data that makes it easier for a more similar approach to do well. Given how brittle many of these methods are, this could make a difference. Would it be possible to instead consider some open-source benchmarks such as VD4RL (Lu et al), which would make it "fair" across different approaches? For example you could use the random or mixed datasets for the embodied and then expert for the demonstration. 3. All of these experiments are within the same single environment, with 3 different reward functions. There is no variability in terms of the dynamics or observation space. To take a tiny step towards a "foundational world model" could you consider some variation, such as using the distracting control suite or varied dynamics in the simulator? This is also available in VD4RL so could be used there too. My guess is the world model approach would actually do better and it would make the results more interesting. 4. This is a very active area of research so the citations seem light, for example: - Edwards et al. "Imitating latent policies from observation". ICML 2019 - Seo et al. "Reinforcement learning with action-free pre-training from videos". ICML 2022 - Schmeckpeper et al. "Reinforcement learning with videos: Combining offline observations with interaction". CoRL 2020

Questions

How was the IDM baseline tuned? For example, in the VPT paper they mention different architecture choices made a big difference for performance.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Limitations are discussed thoroughly in the main body.

Summary

This paper presents an imitation learning approach by first training a world model to predict next observations conditioned on (given) actions, and in a second phase training a policy that amortizes action inference by maximizing the likelihood of observations under a dataset of expert demonstrations. The authors compare their method to BCO(0) on the DMC walker and cheetah environments.

Strengths

The paper is well written, the idea is clearly explained and the benchmark seems well executed.

Weaknesses

The experimental results are hard to parse and have some anomalies (see my questions). Also the strongest claims "AIME outperforms the baselines by a large margin, which indicates the strong generalisability of a forward model over an inverse model. We also find that AIME makes substantially better use of exploratory data." are also mainly based on the Walker experiment, but are less outspoken on the Cheetah dataset.

Questions

How is the performance measured? Is this accumulated reward (normalized to the expert performance)? There are some questioning data points in the evaluations, for instance: - AIME on Cheetah run->flip is able to learn from Visual, but not at all from LPOMDP/MDP, same for flip->run - On flipb->run on Cheetah is impossible to learn, but running backwards is - BCO(0)-MDP seems to outperform AIME-MDP on xxx->flipb Do you have any insights on these results, whether this is due to the task at hand, the collected dataset, ...? Moreover, on Cheetah the resulting agent reaches an overall performance of under 50%. Any idea why this performance gap is there?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

The authors adequately address the limitations in the conclusion.

Summary

This paper presents an algorithm named AIME to learn the world model and apply it to downstream tasks. In the first stage, AIME learns a world model from a dataset with actions to maximize the likelihood via EBLO. While in the second stage, given observation-only demonstrations, AIME optimizes the action sequence to imitate the expert’s behavior. The empirical result shows that AIME outperforms previous methods in DMC tasks.

Strengths

1. Optimizing actions from observation-only trajectory via ELBO is somewhat novel compared to behavior cloning methods. 2. AIME performs better than previous methods even with local and image-based observations.

Weaknesses

1. The major concern is the problem setting of AIME. In my view, more general setting the agent can only get state trajectory in the world-model learning stage while can obtain action-labeled data in the second stage. Then the agent has much more data in training (e.g., human data without actions) and only need a small amount of action-labeled data for fast adaptation. The authors should clarify the significance of the problem settings studied in this paper. 2. Since the world model applies action-labeled data in the first stage, the world model will be related to the policies that generate the dataset. I wonder if the model can handle datasets with a mixture of policies or low-quality policies.

Questions

Why the performance seems to have a very large variance in Figure 3?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

N/A

Summary

The paper proposes action inference by maximising evidence as a way for an MBRL to replicate most likely actions using appropriate world models. The algorithm has two phases: 1) Learn the world model based on a replay buffer, and 2) imitate the expert's behaviour by inferring the policy that maximizes the evidence of the demonstration under the policy and world model. Experimental results on the Walker and Cheetah embodiments of the DeepMind Control Suite demonstrate that this zero-shot imitation performance outperforms the current state-of-the-art approaches.

Strengths

- The paper addresses a major issue in deep reinforcement learning (DRL), namely sample inefficiency. By suggesting a method that can harness observational data, the authors propose a way to improve the sample efficiency of DRL agents. - The paper introduces a new method, Action Inference by Maximising Evidence (AIME), for imitation learning. This method is designed to mimic the human ability to learn quickly from observation, which is an interesting contribution. - AIME's two-phase learning process, involving the creation of a world model and then using it for imitation learning, is a unique approach. This process allows the agent to understand its own body and the likely actions that led to observed behaviors, which is a crucial aspect of learning. - The method is capable of "zero-shot" learning, meaning it does not require further training for the world model or online interactions with the environment after being given the demonstration. This is a significant advantage in terms of efficiency and practicality. - The authors provide empirical validation of their method on the Walker and Cheetah embodiments of the DeepMind Control Suite. They demonstrate that their method outperforms state-of-the-art baselines, which strengthens the credibility of their approach.

Weaknesses

- The paper validates the AIME method using the Walker and Cheetah embodiments of the DeepMind Control Suite, which are simulated environments. It's unclear how well the method would perform in real-world scenarios, where conditions can be more complex and unpredictable. - The AIME method assumes that the agent can learn a perfect world model from its past experiences. This may not always be possible due to the complexity and unpredictability of many environments. The performance of the method could be affected if the world model is not accurate. - The effectiveness of the AIME method may heavily depend on the quality of the observation-only demonstrations provided. If these demonstrations are not representative of the task at hand, or if they are of poor quality, the performance of the method could be significantly affected. - The paper does not discuss the computational complexity of the AIME method or its scalability to larger and more complex tasks. If the method is computationally intensive, it may not be practical for use in real-time applications or on larger scales.

Questions

- Does the world model training in Phase 1 have to converge because imitation learning can happen? Is this primarily for changes in the task? E.g., going from walking to hopping but with the same agent in the same environment? - What if the space of actions change between phase 1 and phase 2? Will AIME still work? It doesn’t seem like it. - Can action trajectories be learnt directly instead of one step action policies? - How sensitive is the AIME method to the quality and diversity of the observation-only demonstrations provided? What happens if the demonstrations are not representative of the task at hand or are of poor quality? - What is the computational complexity of the AIME method? How well does it scale to larger and more complex tasks? - does the AIME method handle situations where the world model learned from past experiences is not accurate or complete? What are the implications if the world model is imperfect? - How well does the AIME method perform in terms of transfer learning? Can the world model learned in one context be effectively applied to another for the policy? - How robust is the AIME method to changes in the environment or task? Can it adapt to new situations without requiring additional training?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

The authors point out the following limitations that limit the scope of the current results: - The AIME method performs well with visual input, but there is a significant performance gap when compared to the LPOMDP setting where low-dimensional signals are observed. This is attributed to the loss surface of the pixel reconstruction loss not being smooth enough to allow the gradient method to find an equally good solution. - The study only considers the simplest setting where both the embodiment and sensor layout are fixed across tasks. This is a limitation as humans observe others in a third-person perspective and can imitate animals whose bodies are not similar to humans. Relaxing these assumptions could allow for transfer across different embodiments and even directly from human videos. - For some tasks, even humans cannot achieve zero-shot imitation by only watching others. This could be due to the task's complexity or completely unfamiliar skills. Even with proper instruction, humans still need to practice in the environment and learn something new to solve some tasks. This suggests the need for an online learning phase as an extension to the AIME framework.

Seems good!

Thank you for your response, it seems sensible. I would still like to see something a bit more ambitious for a super high score, but I think this paper should be accepted. The new experiments do provide incremental confidence given they are open source benchmarks vs. author generated settings. Given the other scores are borderline I am willing to support it with a 7.

Response

Thanks for the response. The limitation of this work in low-quality data needs more research in the future. I will keep the score unchanged.