Zero-shot Policy Learning with Spatial Temporal RewardDecomposition on Contingency-aware Observation

TL;DR

This paper introduces a zero-shot policy learning method that decomposes sparse rewards into contingency-aware observations, enabling generalization to unseen environments without additional training.

Contribution

The proposed approach decomposes sparse rewards into finer-grained, contingency-aware observations and uses these to enable zero-shot generalization via MPC in new environments.

Findings

Outperforms behavior cloning and state-of-the-art RL on benchmark tasks

Effective in both video game and robotic control environments

Demonstrates strong zero-shot transfer capabilities

Abstract

It is a long-standing challenge to enable an intelligent agent to learn in one environment and generalize to an unseen environment without further data collection and finetuning. In this paper, we consider a zero shot generalization problem setup that complies with biological intelligent agents' learning and generalization processes. The agent is first presented with previous experiences in the training environment, along with task description in the form of trajectory-level sparse rewards. Later when it is placed in the new testing environment, it is asked to perform the task without any interaction with the testing environment. We find this setting natural for biological creatures and at the same time, challenging for previous methods. Behavior cloning, state-of-art RL along with other zero-shot learning methods perform poorly on this benchmark. Given a set of experiences in the training environment, our method learns a neural function that decomposes the sparse reward into particular regions in a contingency-aware observation as a per step reward. Based on such decomposed rewards, we further learn a dynamics model and use Model Predictive Control (MPC) to obtain a policy. Since the rewards are decomposed to finer-granularity observations, they are naturally generalizable to new environments that are composed of similar basic elements. We demonstrate our method on a wide range of environments, including a classic video game -- Super Mario Bros, as well as a robotic continuous control task. Please refer to the project page for more visualized results.