Generalizable Imitation Learning from Observation via Inferring Goal Proximity

1 Introduction . Humans are effective at learning a task from demonstrations and applying the learned behaviors to other situations . We achieve this by extracting the underlying structure of the task when observing others fulfilling the task , instead of simply memorizing the demonstrator ’ s low-level actions [ 4 , 18 ] . This high-level task structure generalizes to new situations and thus helps us to quickly learn the task in new situations . One intuitive and readily available instance of such high-level task structure is task progress , measuring how much of the task the agent completed . Inspired by this insight , we propose a novel imitation learning method that utilizes task progress for better generalization to unseen states and goals . Typical learning from demonstration ( LfD ) approaches [ 13 , 35 ] greedily imitate the expert policy and thus suffer from accumulated errors causing a drift away from states seen in the demonstrations [ 38 ] . To make the imitation policy more robust to states not in demonstrations , adversarial imitation learning methods [ 14 , 17 ] encourage the agent to stay near the expert trajectories using a learned reward that distinguishes expert and agent behaviors . However , such learned reward functions often overfit to the expert demonstrations by learning spurious correlations between task-irrelevant features and expert/agent labels [ 52 ] , and thus suffer from generalization to slightly different initial and goal configurations from the ones seen in the demonstrations ( e.g . holdout goal regions or larger perturbation in goal sampling ) . To learn a more generalizable and informative reward from demonstrations , we propose an imitation learning from observation ( LfO ) method , which learns a task progress estimator and uses the task progress estimate as a dense reward for training a policy as illustrated in Figure 1 . Unlike discriminating expert and agent behaviors by predicting binary labels in prior adversarial imitation ∗Equal contribution . Correspondence to : lee504 @ usc.edu and aszot3 @ gatech.edu †This work was partially carried out during an internship at NAVER AI Lab . ‡AI Advisor at NAVER AI Lab . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . learning methods , which is prone to overfitting to task-irrelevant features , the task progress estimator is required to learn more task-relevant information to precisely predict the task progress on a continuous scale . Hence , it can generalize better to unseen states and provide more informative rewards . As a measure of progress in goal-directed tasks , we define goal proximity , which is an estimate of temporal distance to the goal ( i.e . the number of actions required to reach the goal ) and entails all semantic information about how to reach the goal . We then train a proximity function to predict the goal proximity from expert demonstrations and agent experience . This proximity function acts as a dense reward to guide a reinforcement learning agent to reach states with high proximity , leading to the goal . In this paper , we focus on learning the proximity function and policy in a state space shared by the expert and learner , and leave generalizing to different embodiments as future work . However , the predicted goal proximity can still be inaccurate on states not in the demonstrations , resulting in unstable policy learning . To improve the accuracy of the proximity function , we continually update it with trajectories from both the expert and learning agent . In addition , we penalize trajectories with the uncertainty of the proximity prediction to prevent the policy from exploiting inaccurate high proximity predictions . By leveraging the agent experience and predicting proximity function uncertainty , the proposed method achieves more efficient and stable policy learning . The main contribution of this paper is an LfO algorithm for goal-directed tasks with better generalization to new goals or states not in demonstrations using goal proximity that informs an agent of the task progress . Together with a difference-based reward and uncertainty penalty of goal proximity estimation , our method provides more informative and robust rewards . Our extensive experiments show that the policy learned with the goal proximity function generalizes better than the state-ofthe-art LfO algorithms on various goal-directed tasks , including navigation , locomotion , and robotic manipulation . Moreover , our method shows comparable results with LfD methods which learn from expert actions and a goal-conditioned imitation learning method which uses a sparse task reward . 2 Related Work . Imitation learning [ 39 ] aims to leverage expert demonstrations to acquire skills . While behavioral cloning [ 35 ] is simple but effective with a large number of demonstrations , it suffers from compounding errors caused by covariate shift [ 38 ] . On the other hand , inverse reinforcement learning ( IRL ) [ 1 , 29 , 51 ] estimates the underlying reward from demonstrations and trains a policy through reinforcement learning ( RL ) with this reward , which can better handle the compounding errors . Specifically , generative adversarial imitation learning ( GAIL ) [ 17 ] shows improved demonstration efficiency by training a discriminator to distinguish expert and agent transitions and using the discriminator output as a reward for policy training . GoalGAIL [ 9 ] further improves sample efficiency for goal-directed tasks by relabeling transitions [ 2 ] and using true environment rewards . While these imitation learning algorithms require expert actions , imitation learning from observation ( LfO ) approaches learn from state-only demonstrations , such as videos and kinesthetic demonstrations . To imitate demonstrations without expert actions , inverse dynamics models [ 30 , 33 , 46 ] , reachability functions [ 23 ] , or learned reward functions [ 11 , 26 , 41 , 42 ] can be learned and used for policy training , but training such models requires a large amount of quality data or additional test-time demonstrations . On the other hand , state-only adversarial imitation learning [ 47 ] can imitate from a few demonstrations . However , in such adversarial imitation learning approaches , the discriminator tends to find spurious associations between task-irrelevant features and expert/agent labels [ 52 ] . This becomes problematic when the agent encounters unseen states and the discriminator erroneously assigns agent behaviors low scores based on these task-irrelevant features , providing a poor reward for the agent . To overcome finding spurious associations , in addition to discriminating expert and agent trajectories , we propose to also estimate the proximity to the goal , which requires more task-relevant information and thus generalizes better to new states . Temporal progress estimation has shown its effectiveness as an auxiliary reward for RL [ 10 , 24 , 27 ] and decision making criteria [ 3 , 6 , 8 ] . However , these methods learn the progress estimator only from the given demonstrations . This hinders policy learning when the progress estimator fails to generalize to agent experience , allowing the agent to exploit inaccurate progress predictions for higher reward . Moreover , greedily choosing an action with the highest predicted temporal progress [ 3 , 6 , 8 ] could lead to low long-term returns . By incorporating online updates , uncertainty estimates , and a difference-based proximity reward , our method robustly learns from demonstrations to solve goal-directed tasks without access to expert actions or the true environment reward . 3 Method . In this paper , we address the problem of LfO for goal-directed tasks with a focus on generalization to states or goals not covered in the demonstrations . Adversarial LfO methods [ 47 , 49 ] suggest learning a reward function that penalizes agent state transitions deviating from the expert trajectories . However , these learned reward functions often focus on task-irrelevant features [ 52 ] and do not generalize to states not in the demonstrations , leading to unsuccessful policy training . To learn a generalizable reward , we propose to leverage task progress information freely available in demonstrations , in terms of goal proximity , which estimates temporal distance to the goal ( i.e . number of actions required to reach the goal ) . Predicting precise goal proximity on a continuous scale , rather than simply distinguishing expert and agent states , requires the model to capture task-relevant information , allowing the proximity prediction to generalize to states not in the demonstrations ( Section 3.2 ) . Then , a policy learns to reach states with higher proximity prediction , leading to the goal ( Section 3.3 ) . Moreover , we propose to use the uncertainty of the proximity prediction to prevent the policy from exploiting over-optimistic proximity predictions and yielding undesired behaviors . 3.1 Preliminaries . We formulate our problem as a Markov decision process [ 44 ] defined through a tuple ( S , A , R , P , ρ0 , γ ) of the state space S , action space A , reward function R ( s , a , s′ ) , transition distribution P ( s′|s , a ) , initial state distribution ρ0 , and discounting factor γ . We define a policy π ( a|s ) that maps from a state s to an action a and correspondingly moves an agent to a new state s′ according to the transition probability P ( s′|s , a ) . The policy is trained to maximize the sum of discounted rewards , E ( s0 , a0 , ... , sTi ) ∼π [ ∑Ti−1 t=0 γ tR ( st , at , st+1 ) ] , where Ti is the variable episode length . In imitation learning , the learner receives a set of N expert demonstrations , De = { τe1 , . . . , τeN } . In this paper , we specifically consider the LfO setup where each demonstration τei is a sequence of states . Moreover , we assume that goal information is explicitly or implicitly included in the state s , and all demonstrations are successful ; therefore , the final state of each trajectory achieves the task goal . 3.2 Learning Goal Proximity Function . To effectively leverage expert demonstrations and generalize to new states or new goals , learning a generalizable reward function is essential . In goal-directed tasks , an estimate of how close an agent is to the goal can be utilized as a dense and direct learning signal . Moreover , predicting the continuous goal proximity requires understanding the task structure and thus encourages finding more task-relevant features , resulting in better generalization . Therefore , instead of learning to simply discriminate agent and expert trajectories , we propose to learn a goal proximity function , f : S → [ 0 , 1 ] , which predicts goal proximity of a state s , which is a discounted value based on the temporal distance to the goal ( i.e . inversely proportional to the number of actions required to reach the goal ) . In this paper , we define goal proximity as the exponentially discounted proximity f ( st ) = δ ( Ti−t ) , where δ ∈ ( 0 , 1 ) is a discounting factor and Ti is the episode length . Note that the goal proximity function measures the temporal distance , not the spatial distance , between the current and goal states . Therefore , a single proximity value can entail all information about the task , goal , and any roadblocks . There are alternative ways to define goal proximity , such as linearly discounted proximity [ 24 ] and ranking-based proximity [ 5 , 6 ] . But , in this paper , we use the exponentially discounted proximity as it performs better across most tasks ( see appendix , Figure 8 ) . We train a goal proximity function fφ parameterized by φ to minimize the following objective : Lφ = Eτei ∼De , st∼τei [ fφ ( st ) − δ ( Ti−t ) ] 2 . ( 1 ) Since the goal proximity function trained only on expert demonstrations can overfit to the data , we further train the goal proximity function with online agent experience by setting the target proximity of states in agent trajectories to 0 , similar to adversarial imitation learning methods [ 17 ] : Lφ = Eτei ∼De , st∼τei [ fφ ( st ) − δ ( Ti−t ) ] 2 + Eτ∼πθ , st∼τ [ fφ ( st ) ] 2 . ( 2 ) By learning to predict the goal proximity , fφ not only learns to discriminate agent and expert trajectories ( i.e . predict 0 proximity for an agent trajectory and positive proximity for an expert trajectory ) but also acquires the task information about temporal progress entailed in the trajectories . From this freely available additional supervision , the proximity function is required to learn taskrelevant features . Hence , the resulting proximity function generalizes better to unseen states and provides more informative learning signals to the policy as empirically shown in Section 4 . Due to the lack of environment reward , successful agent experience is also used as negative examples for proximity function training , and thus the proximity function learns to predict low goal proximity even for successful trajectories . However , early stopping and learning rate decay can ease this problem [ 52 ] , and the optimal proximity function still outputs the average of expert and agent labels , which is δ ( Ti−t ) /2 for ours and 0.5 for GAIL [ 14 ] .

Summary: Disclaimer:  I was the reviewer of this paper's older version submitted to ICML 2021. Along the reviewing process, the authors' information was kept anonymous.   Overall I'm glad to see this improved version. As the authors stated, 1) comparison to Goal-GAIl added; 2) clearly specify the generalization is limited, and not referring to out-of-demo goals/states which including a different embodiment (for example unseen objects) of environments. Thanks for the authors' improved clarity.  This paper is well written with sound evaluation results, falling in approaches learning reward function from human demonstrations, though I prefer viewing the generalization problem of adapting human demonstrations to unseen states or actions as more of a representation learning or meta-learning problem since such problem settings will flexibly apply to the manipulation of unseen objects, which is important to learn generalizable manipulation skills.