Inversely Learning Transferable Rewards via Abstracted States

TL;DR

The paper introduces TraIRL, a transferable IRL framework that learns a cross-task abstract state space and an abstract, state-only reward from demonstrations across multiple source tasks. By coupling a shared encoder with task-specific decoders (multi-head VAE) and a Wasserstein-based discriminator, TraIRL captures invariant, optimality-relevant structure that transfers to unseen target tasks without demonstrations. The authors provide an analytical framework for reward transferability, validate the approach on MuJoCo-Gym and AssistiveGym, and show that abstract rewards yield superior transfer performance compared to baselines, including cross-domain transfers. The work demonstrates that disentangling reward structure from dynamics via abstraction can markedly improve generalization in robotics and human-robot interaction settings. Practical implications include improved plug-and-play reward shaping for new tasks and more robust policy learning in varied environments.

Abstract

Inverse reinforcement learning (IRL) has progressed significantly toward accurately learning the underlying rewards in both discrete and continuous domains from behavior data. The next advance is to learn {\em intrinsic} preferences in ways that produce useful behavior in settings or tasks which are different but aligned with the observed ones. In the context of robotic applications, this helps integrate robots into processing lines involving new tasks (with shared intrinsic preferences) without programming from scratch. We introduce a method to inversely learn an abstract reward function from behavior trajectories in two or more differing instances of a domain. The abstract reward function is then used to learn task behavior in another separate instance of the domain. This step offers evidence of its transferability and validates its correctness. We evaluate the method on trajectories in tasks from multiple domains in OpenAI's Gym testbed and AssistiveGym and show that the learned abstract reward functions can successfully learn task behaviors in instances of the respective domains, which have not been seen previously.

Figures (10)

• Figure 1: TraIRL framework overview. Expert and learner trajectories in multiple source tasks are mapped to shared abstract states via a shared encoder. A discriminator compares the abstracted state densities to estimate the 1-Wasserstein distance between expert and learner, which guides learning a reward function over the abstract space through a covariance-based objective. The learned reward function then optimizes the learner policy.
• Figure 2: Source and target tasks from MuJoCo-Gym domains. Red legs are the disabled legs of the robots. Frames (a,b) depict source tasks of running with a disabled leg in Half Cheetah while (c) represents the target environment with no disability. Similarly, frames (d,e) show the source tasks of running with different pairs of disabled legs in Ant, whereas (f) shows the target of running with another pair of disabled legs.
• Figure 3: Visualization by t-SNE of (a) sampled abstractions and (b) ground states for Half Cheetah, and analogously for Ant (c,d) source environments. Details of these visualizations are in Appendix \ref{['appendix:tsne']}
• Figure 4: Two tasks in the feeding domain (a) are used as sources to learn a reward function that is transferred to perform the task of scratching an itch (b). (c) Cumulative reward with standard deviation in the Assistive Gym environments. TraIRL has the highest reward in the target task.
• Figure 5: Smoothed training curve for Half Cheetah in two source tasks. AIRL-ME and $f$-IRL perform poorly in the experiments and are therefore excluded from the comparison.

Theorems & Definitions (14)

• Definition 1: Cross-Task Abstraction
• Theorem 1: Gradient of Reward Function
• Definition 2: Reward Transferability
• Theorem 2: Applicability of TraIRL
• proof : Sketch of Proof
• proof
• proof
• Definition 3: Induced ground-level reward function
• Definition 4: Disentangled rewards
• Definition 5: Decomposability condition