[Paper Review] Learning Representations that Enable Generalization in Assistive Tasks

Jerry Zhi-Yang He, Aditi Raghunathan+3 more|arXiv (Cornell University)|Dec 5, 2022

Reinforcement Learning in Robotics4 citations

TL;DR

This paper proposes PALM (Prediction-based Assistive Latent eMbedding), a framework that learns a disentangled latent representation of human partner policies directly from interaction history to improve zero-shot generalization in assistive robotics. By training the latent space to predict human actions and enabling test-time adaptation using observed behavior, PALM achieves superior generalization to out-of-distribution human policies compared to state-of-the-art sim2real and human-robot interaction methods.

ABSTRACT

Recent work in sim2real has successfully enabled robots to act in physical environments by training in simulation with a diverse ''population'' of environments (i.e. domain randomization). In this work, we focus on enabling generalization in assistive tasks: tasks in which the robot is acting to assist a user (e.g. helping someone with motor impairments with bathing or with scratching an itch). Such tasks are particularly interesting relative to prior sim2real successes because the environment now contains a human who is also acting. This complicates the problem because the diversity of human users (instead of merely physical environment parameters) is more difficult to capture in a population, thus increasing the likelihood of encountering out-of-distribution (OOD) human policies at test time. We advocate that generalization to such OOD policies benefits from (1) learning a good latent representation for human policies that test-time humans can accurately be mapped to, and (2) making that representation adaptable with test-time interaction data, instead of relying on it to perfectly capture the space of human policies based on the simulated population only. We study how to best learn such a representation by evaluating on purposefully constructed OOD test policies. We find that sim2real methods that encode environment (or population) parameters and work well in tasks that robots do in isolation, do not work well in assistance. In assistance, it seems crucial to train the representation based on the history of interaction directly, because that is what the robot will have access to at test time. Further, training these representations to then predict human actions not only gives them better structure, but also enables them to be fine-tuned at test-time, when the robot observes the partner act. https://adaptive-caregiver.github.io.

Motivation & Objective

• To address the challenge of generalization in assistive robotics where test-time human partners may be out-of-distribution (OOD) relative to the training population.
• To design a framework that learns a structured, adaptable latent space of human policies based on interaction history, rather than environment parameters or indirect signals like rewards.
• To enable test-time adaptation of the latent space using real interaction data, improving robustness to OOD human behaviors.
• To evaluate and compare representation learning methods in a controlled, OOD-focused assistive task setting.

Proposed method

• The framework jointly trains a robot policy and a latent embedding space using end-to-end backpropagation, with the latent space trained to predict human actions from interaction history.
• The latent space is optimized via a contrastive objective and a KL regularization term to encourage smooth, generalizable structure.
• Test-time adaptation is enabled by fine-tuning the latent space using observed human actions during interaction, allowing dynamic alignment with new partners.
• The method uses a two-stream architecture: one for encoding human history and one for predicting next actions, with the latent code serving as a shared representation.
• The framework is evaluated in two simulated assistive tasks: a 2D reach environment and a 2D itchy-scratch task with diverse human policies.
• Human policies are explicitly constructed to be OOD, enabling controlled evaluation of generalization performance.

Experimental results

Research questions

• RQ1 Can a latent representation of human policies trained on interaction history generalize better to OOD human behaviors than representations based on environment parameters or indirect signals?
• RQ2 Does predicting human actions during training lead to a more structured and generalizable latent space compared to predicting states or rewards?
• RQ3 To what extent does test-time adaptation of the learned latent space improve performance on unseen human partners?
• RQ4 How does the choice of representation learning objective (e.g., action prediction vs. reward prediction) affect generalization in assistive tasks?

Key findings

• PALM with action prediction outperforms all baselines, including RMA, RNN, LILI, and RILI, in generalization to OOD human policies across both the reach and itch-scratch tasks.
• The latent space learned via action prediction successfully captures the underlying ring-like structure of human policies, while RMA and RILI fail to preserve this structure due to reliance on ambiguous or non-informative signals.
• KL regularization improves generalization by inducing a smoother latent distribution, allowing better interpolation of OOD humans into the latent space, especially in missing regions of the distribution.
• Test-time adaptation significantly improves performance, with the latent space of PALM better embedding OOD humans by filling in gaps in the distribution, as shown in visualization of the latent space.
• Baselines that predict rewards or states fail to generalize well, as these signals do not capture the structural relationships between human policies.
• The method shows strong performance even when the human distribution becomes more complex, outperforming MLP and RILI, which degrade under increased distributional complexity.