Learning Representations that Enable Generalization in Assistive Tasks
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 di  cult 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 beneﬁts 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 ﬁnd 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 ﬁne-tuned at test-time, when the robot observes the partner act.
https://adaptive-caregiver.github.io
Introduction
Our ultimate goal is to enable robots to assist people with day to day tasks.
In the context of patients with motor impairments, this might mean assistance with scratching an itch, bathing, or dressing [1, 2, 3].
These are tasks in which doing reinforcement learning from scratch in the real world is not feasible, and so sim2real transfer is an appealing avenue of research.
Sim2real methods for physical robot tasks in isolation typically work by constructing a diverse "population" of environments and training policies that can work with any member of the population (e.g. a range of parameters of a physics simulator or a range of lighting and textures) [4, 5, 6, 7, 8, 9, 10, 11].
Similarly, population-based (self-play) training has proven successful in zero-sum games against humans [12, 13, 14].
But unlike tasks the robot does in isolation, assistance requires coordinating with a human who is also acting.
And unlike competitive settings, assuming the human to be optimal when they are not, can result in dramatically poor performance [15].
Thus, in sim2real for assistance, we have to design a population of potential users and strategies to train with, akin to the physical environment parameters in typical sim2real tasks, rather than the standard population-based training approaches used in competitive settings.
But designing a population that is diverse and useful enough to enable generalization to test-time humans, each with their own preferences, strategies, and capabilities, remains very challenging, making it likely that test-time partners might lie outside of the distribution the population was drawn from.
Therefore, sim2real methods for assistance will need to be ready to generalize to out-of-distribution partner policies.
In this work, we identify two principles as key to enabling better generalization.
First is that we beneﬁt from learning a latent space of partners that distills their policies down to a structure that is useful for the robot’s policy and that makes it easy to identify partners at test time.
Second is that we need to be prepared for this space to not perfectly capture the space of real human policies, and design it so that it is adaptable with real test-time interaction data.
We thus propose a framework that learns a latent space directly from history of interaction by predicting the partner’s actions .
Our framework allows a robot to capture the relevant information about the human partner that the robot can actually identify when starting to interact, and also enables test-time adaptation of the latent space itself when observing the partner’s actions.
When evaluated with partner policies we purposefully design to be out-of-distribution, we ﬁnd that our approach leads to better generalization than prior methods which either do not learn a latent space at all [16], do not learn a latent space directly based on interaction history [7], or train a latent space based on other observables, like states or rewards [17, 18].
Our contributions are four-fold:
We introduce an assistive problem setting where the focus is explicitly on generalization to out-of-distribution partner policies.
We introduce a framework for training policies for this problem setting,
Prediction-based Assistive Latent eMbedding (PALM).
This enables us to study di  erent methods for learning latent representations on how well they enable generalization.
We identify that the design choice of training a latent space by predicting partner actions directly from history outperforms (1) state-of-the-art sim2real approaches used in non-assistive tasks that are based on embedding environment parameters [7] as well as (2) human-robot interaction approaches that train representations by predicting observed states or rewards [17, 18].
We propose to adapt the learned latent space at test time, upon observing the partner’s actions, and show it leads to generalization performance gains.
The Assistive Personalization Problem
In this section, we introduce the personalization problem in an assistive context.
In particular, our goal is to learn a robot policy c ' that can assist a novel human partner in zero-shot fashion, or with a small amount of test-time data.
Two-player Dec-POMDP .
An assistive task can be modeled as a two-agent, ﬁnite horizon decentralized partially-observable Markov decision process (Dec-POMDP) and is deﬁned by a tuple h ( , U ,  ' ,   , T , ⌦ ' , ⌦  , $ , ' i .
Here ( is the state space and  ' ,   are the human’s and the robot’s action spaces, respectively.
The human and the robot share a real-valued reward function ' : ( ⇥  ' ⇥   !
R ; however, we assume that the reward function is not necessarily observed by the robot, i.e. its parameters (e.g. the location of the human has an itch) are in the hidden part of the state.
T : ( ⇥  ' ⇥   ⇥ ( !
[ 0 , 1 ] is the transition function, which outputs the probability of the next state given the current state and all agents’ actions.
⌦ ' and ⌦  are the sets of observations for the robot and human, respectively, and $ : ( ⇥  ' ⇥   !
⌦ ' ⇥ ⌦  represents the observation probabilities.
We denote the horizon of the MDP by ) .
Target User.
We target users with partial motor functions — a common impairment for individuals with partial arm functions.
This is an impairment that can occur in some people with cervical SCI,
ALS, MS, and some neurodegenerative diseases — leading to the need for robotic assistance.
We model the extent of the impairment as the privileged information in the Dec-POMDP.
The robot does not know this a-priori and thus needs to adapt to individual users’ capabilities.
The Robotic Caregiving Setup.
We deﬁne the observation space for the robot and the human following [3]: the robot observes its own joint angles, and the human’s joint positions in the world coordinate and contact forces; the human observes their joint angles (proprioception) and the ende  ector position of the robot.
When training with simulated humans, the robot gets a reward signal (which depends on privileged information), and has to use that signal to learn to implicitly identify enough about the human to be useful; at test time, the robot does not observe reward signal and must use what it has learned at training time to identify the human’s privileged information and be helpful.
Distributions of Humans.
Let function c
: ⌦ ⇤  ⇥   !
[ 0 , 1 ] be the human policy that maps from local histories of observations o C =
( >  1 ,..., > C ) over ⌦  to actions.
We deﬁne two distributions of human policies c
2 D train , D test .
In the assistive itch scratching, D train can be a set of humans with di  erent itch positions on their arms, which lead to their di  erent movements.
We refer to them as in-distribution humans .
D test contains out-of-distribution humans whose itch position di  er from those in the D train .
At training time, the robot has access to D train .
Thus, it has ground-truth knowledge about the training human’s privileged information, such as each human’s itch position.
At test time, we evaluate the robot policy by sampling humans c
⇠ D test and directly pairing them with the robot policy.
We evaluate the zero-shot and few-shot adaptation performance of the robot policy.
Objective.
The main problem we study is how to leverage the training distribution to learn a robot policy c ' : ⌦ ⇤ ' !
' such that we achieve the best performance on test humans.
Concretely, we deﬁne the performance of the robot and human as
[Equation]
Only given access to D train , our objective is to ﬁnd the robot policy c ' = argmax c
( c , c  ) , c  ⇠ D test .
Learning Personalized Embeddings for Assistance with PALM
In this section, we present Prediction-based Assistive Latent eMbedding (PALM).
We introduce the general framework of using a latent space to perform personalization in an assistive context.
We then highlight the advantage of action prediction in contrast to prior works.
Finally, we describe how we can optimize PALM at test time to work with out-of-distribution humans.
3.1 Learning an Assistive Latent Space
Given a training distribution of humans D train 1 , we would like to learn a robot policy that can adapt to assist new users.
To achieve that, a robot must learn to solve the task while e  ciently inferring the hidden component that di  ers across humans.
One natural way to do so is to learn a latent space that succinctly captures what di  ers across humans in a way that a  ects the robot’s policy.
When deployed on a test human, the robot infers this latent embedding and uses it to generate personalized assistance.
We denote the latent space as I C ⇠ E \ ( I ; g 1: C ) , where E \ encodes the trajectory g so far and outputs latent vector
I
C .
The robot uses this latent space to compute its actions 0 'C
= c ' ( > 'C , I C ) .
We train the base policy c ' and the latent space encoder E \ jointly as they are interdependent [19] — better robot policy leads to di  erent trajectories across humans, which in turn leads to distinguishing I .
See Appendix D on training details.
Ideally, we would like I to capture su  cient information to di  erentiate the humans, similar to performing a “dimensionality reduction” on human policies c  .
We hereby introduce di  erent objectives for learning such latent space, and how our method — learning by action prediction — makes a good ﬁt for the assistive personalization problem.
3.2 How to Construct the Latent Space
Prior work and limitations LILI [17] and RILI [18] learn a latent embedding of the humans by predicting the next observations and rewards.
While they have been shown to work in predicting and inﬂuencing human behaviours, both methods assume access to the reward function at test time, which we do not have access to in the assistive setting — we don’t know a-priori the preference and needs of a new user.
RMA [7] enables fast robot adaptation by learning a latent space of environment parameters, such as friction, payload, terrain type, etc.
While it works for a single robot, it is unclear in human-robot settings, how we can encode human motions and preferences as environment parameters.
Learning by action prediction.
Given history g C  # : C =  [Equation] [Equation]  of # robot observation and human action 2 pairs, as outlined in Fig.
1, we embed this trajectory to a lowdimensional manifold and use it to predict 0 C .
The intuition is that if we are able to predict the next action by this human, we extract the su  cient information about the human’s policy c  .
The latent vector I is representative of the trajectory so far and indicative of the person’s future actions.
We do this by training a decoder D q parameterized by q to predict the next action from the encoder’s output I ⇠ E \
( I ; g 1: C ) .
[Equation]
The encoder E is a recurrent neural network parameterized by \ .
Here the second term is a regularization term motivated by Variational Autoencoder [20, 21], that enforces the latent space to follow a normal prior distribution.
This encourages nearby terms in the latent space to encode similar semantic meanings.
In the context of assistive tasks, this helps us better cluster similar humans closer in the latent space, and we show a didactic example in Sec. 4.3.
3.3 Latent Space Adaptation at Test Time
At test time, as we work with a new user, we would like our encoding of the new user to match the true latent information, I ⇤ of that user.
In other words, we would like to minimize ||E( g )  I ⇤
||
Because we do not know about the new users a-priori, we can only optimize for this objective via proxy, which we refer to as test time adaptation .
Since the PALM latent space is based on action prediction, we can adapt it to a new user by further optimizing the latent space.
Note that because Eq.
(2) requires only observation-action data, we do not need any additional label to perform test time adaptation.
More formally, we collect a small dataset of test-time interaction trajectories, g , and perform a few gradient steps to optimize both the encoder and decoder for Eq.
(2): ( \ , q ) !
( \ , q
[Equation]
[Equation] r ( \ , q ) L pred (E \ , D q , g ) .
The idea of test-time optimization has been shown to improve perceptual robustness for grasping in sim2real research [22].
We follow a similar pipeline, where we can improve the latent encoding by collecting unsupervised action data from test users.
Here the main di  erence is instead of perceptual di  erences, our goal is to reduce the domain gap on test users.
Experiments
In this section, we evaluate our method PALM (Prediction-based Assistive Latent eMbedding) in collaborative human-robot environments of varying tasks and varying populations of human models.
In particular, we focus on the out-of-distribution generalization by constructing di  erent forms of out-of-distribution populations.
We focus on empirically investigating the beneﬁts of learning a latent space, the e  ect of di  erent kinds of prediction on learning a useful latent space, the properties of learned latent spaces, and the gains from test-time adaptation to humans.
4.1 Environments
Here we introduce two environments where we study assistive personalization.
In both environments, the robot has to infer some hidden information from the human in order to successfully solve the task.
Note that these are examples meant for demonstrating the e  ectiveness of the algorithm, and we do not claim to solve the full robotic caregiving problem.
Assistive Reacher (Fig. 2) is 2D collaborative environment where a two-link robot arm assists a point human agent to get to the target position.
This target is located at ( 3 cos U H , 3 sin U H ) , where 3 is a ﬁxed value, and U H
[  c , c ] is known to the human, but not the robot.
The human agent is initialized randomly in the 2D plane with random hidden parameters U H
[  c , c ] , : H 2 [ 0 . 5 , 1 . 5 ] .
The robot can only identify the target position by physical interactions — once the robot initiates contact, the human applies a force : H · ( cos U H + c 2 , sin U H + c 2 ) .
Only by recognizing the human in terms of U H , : H can the robot compensate the force, and successfully move the human to the hidden target.
Each episode hastimesteps.
The Scope of Generalization.
We deﬁne D train as 36 samples uniformly sampled from U H[  c , c 2 ] , : H 2 [ 0 . 5 , 1 . 5 ] and D test as 12 samples uniformly sampled from U H 2 [ c 2 , c ] , : H 2 [ 0 . 5 , 1 . 5 ] .
Assistive Itch Scratching (Fig. 1) is adapted from assistive gym [3].
It consists of a human and a wheelchair-mounted 7-dof Jaco robot arm.
The human has limited mobility — they can only move the 10 joints on the right arm and upper chest, and needs the robot’s assistance to scratch the itch.
An itch spot is randomly generated on the human’s right arm.
The robot does not directly observe the itch spot, and relies on interaction with the human to infer its location.
Each episode has 100 timesteps.
We use co-optimization to create D train and D test for Assistive Itch Scratching.
A beneﬁt of the co-optimization framework is that it naturally induces reward-seeking behaviour from the human and the robot, which simulates assistance scenarios.
For instance, to generate more inactive human policies, we can introduce a weighting term in the reward function for human action penalties ' p = 2 p · || c H ( B C )
|| 2 where 2 p is a constant controlling the penalty.
The overall objective becomes
[Equation]
The Scope of Generalization.
We are motivated by real world applications where users tend to have di  erent levels of mobility limitations, or itch locations in di  erent body parts.
To generate a synthetic population to capture such diversity in itch scratching task, we explore di  erent co-optimization settings (1) we assign di  erent human action penalty to be 2 p = 3 , 3 . 5 , 4 , where larger penalties lead to the human agent exerting less e  ort.
We simulate di  erent itch positions on the human’s arm and train co-optimized human and robot policies conditioned on them.
This leads to qualitatively di  erent strategies for the human and the robot.
Note that this serves ﬁrst step to understanding how di  erent methods generalize, since we never expect to be able to capture the diversity in humans perfectly.
For training, we use Proximal Policy Optimization (PPO) to optimize human and robot policies in an interleaving fashion.
Note that we also keep the co-optimized robot policy and use it to obtain expert actions for assistive policy training (see Sec.
3.1 and supplement for full details).
To construct D train and D test , we divide the two arms’ areas into four equal portions, as shown in Fig.
3, and generate human policies conditioned on itch positions in these areas.
D train consists of three of the four portions and D test consists of the remaining one.
We then construct three distribution sets in increasing order of di  culties.
In the ﬁrst distribution D 1 , we conﬁne itch positions from a line-shaped region.
In D 2 , we sample from all the arm areas.
Note that D 1 and D 2 are constructed by setting action penalty 2 p = 3 .
In D 3 , we combine humans of 2 p = 3 , 3 . 5 , 4 , each trained with two di  erent random seeds.
This adds the extra complexity of human activity levels.
We simulatein-distribution humans from each of the three training portions under each action penalty.
4.2 Baselines
We compare with baselines that do not learn an explicit latent space as well as existing methods for adaptation via learning latent embeddings.
MLP and RNN.
We follow [16] that trains sequential models to enable adaptation to simulated humans.
We explore using a recurrent neural network or a feed-forward network on concatenated state-action histories.
The models directly output robot action, and there is no latent space modeling.
ID-based Human Embeddings.
In contrast to learning latent space from history, another class of method studies encodes human-designed environment parameters [7, 19, 23].
We focus speciﬁcally on RMA [7], a two-phased method that ﬁrst learns to encode task-ID (phase I) and then trains a recurrent network to regress to the embeddings from observation history (phase II).
For training a quadruped robot, RMA encodes the physical parameters (friction, payload, etc) of the environment.
The ﬁrst stage trains a policy with ground-truth information, and the second phase performs environment identiﬁcation.
While RMA is shown to be e  ective for learning policy for in-distribution environments, it is unclear how well it generalizes to out-ofdistribution environments.
Furthermore, in assistive tasks, it is unclear how to construct the "ground-truth ID" for phase I that quantiﬁes the user characteristics.
We study
RMA-Param and RMA-Onehot , where we assign each training human a one-hot vector.
For RMA-Param, we use a three-dimensional vector that includes the G , H position of the itch position and the arm index.
When there are multiple human activity levels as mentioned in Sec.
4.1, we introduce a fourth dimension with an integer to indicate the action penalty 2 p .
LILI and RILI.
We consider two other methods of the PALM framework: LILI [17] and RILI [18].
As mentioned in Sec. 3.1, LILI jointly predicts future observation and reward, and RILI predicts reward.
Given that reward is only available at training time, we cannot perform test time optimization for LILI and RILI.
Ablations of PALM.
Our method has several components: we use a recurrent neural network to encode interaction history, and use its output to minimize prediction loss L pred and policy loss L pol .
We also use the KL term in Eq.
(2) to regularize the human embeddings.
To test the e  ectiveness of our method, we separate di  erent parts and create a set of baselines.
We hereby describe them in detail: (1) No L pred : the model shares the same encoder and policy network architecture, yet we don’t optimize for L pred .
By removing the prediction loss in Fig.
1, the latent space is not explicitly trained to contain human information.
(2) No 2 KL : no regularization in the latent space.
(3) Frozen embedding: instead of jointly training embeddings and the policy network, we ﬁrst train the encoder on expert data, freeze it, and then train robot policy.
We include an ablation study of our main experiments in the appendix.
4.3 Didactic Experiment in Assistive Reacher Environment
Can PALM learn a meaningful distribution from the interaction?
Unlike other ID-based methods like RMA, PALM does not have access to the human parameters at training time.
We study whether
PALM can learn a meaningful latent space without explicitly knowing this information.
We sample training humans from U H
[  c , c ] continuously.
We train PALM with di  erent amount of prior regularization,!
from Eq.
We train using a recurrent window of length 4 and a batch size of 512 episodes.
Additional training details can be found in the supplementary material.
We average test results using 100 episodes and visualize the results in Fig. 4.3.
Given that humans are parameterized by U H ,  , the ideal embedding space looks like a ring with a small blob in the center.
The ring corresponds to the 2D projection of U H and the blob denotes the initial part of interaction before contact, which is indistinguishable.
We ﬁnd that while PALM never observes the underlying parameter U H , it can learn a latent space that characterizes U H .
Interestingly, varying the amount of regularization qualitatively a  ects the shape of the latent space.
Setting the VAE regularization 2 KL = 0 .
1 recovers a latent space that most resembles to the ideal latent space.
4.4 Assistive Reacher Main Experiment
Experiment Setting.
We use the ﬁnite D train described in Sec.
and train all baselines forepochs with 512 batch size.
We then evaluate the trained policies on D test .
We normalize the resulting reward with respect to oracle reward.
Results.
We average test results using 100 episodes.
On in-distribution humans, we ﬁnd that all methods successfully follow the right policy that assists the human to reach their goal.
This shows that they all successfully predict the human latent information explicitly or implicitly.
On out-of-distribution humans, the methods are no longer guaranteed to predict the correct embedding.
PALM with action prediction signiﬁcantly outperforms other methods.
With test-time adaptation,
PALM further improves.
Visualizing the latent space.
We qualitatively study generalization by visualizing the latent space as well as the mapped embeddings of both in-distribution and out-of-distribution humans (in red crosses) in Fig.
Interestingly, only PALM with action prediction can infer the “ring” structure.
RMA, RNN,
LILI and RILI fail to do so.
We hypothesize that because hand-crafted human IDs do not convey the information about human policy, RMA warped the IDs in arbitrary what that are harmful for generalization.
The same happens with RILI and LILI.
We hypothesize this is due to the inherent ambiguity in reward prediction: a low reward does not necessarily recover the human policy structure.
The visualization also o  ers some insight into why having 2 KL regularization is helpful for generalization.
Compare the latent space of PALM 2 KL = 0 and PALM 2 KL = 0 .
1 , the latter induces a smoother distribution where test humans are better ﬁtting in the “missing arc” of the “ring".
Further more, we see that with test time optimization, the PALM latent space embeds the OOD human better, by ﬁlling in more of the arc.
4.5 Assistive Itch Scratch Main Experiment
Results.
We follow a similar procedure as the reach environment to train itch scratching policy, and train for 240 epochs with 192 trajectories for batch size.
As shown in Fig.
6, we observe PALM with action prediction has better generalization performance than other baselines.
We see that in the simpliﬁed distribution D 1 , RILI and MLP have the best generalization performance among baselines, yet as the complexity of the training human distribution increases, they deteriorate.
Detailed results of the ablation study are included in the appendix.
Visualizing the Latent Space To further investigate why PALM generalize better to OOD human than
RMA baselines, We visualize the latent space of the "straight-line" distribution.
As we see in Fig.
PALM can capture the structure in human training distribution as two clusters, and also correctly embed the OOD humans distribution as a part of the upper arm distribution.
RMA-based methods, on the other hand, can discover the structure of training humans.
Yet qualitatively, they fail the correctly embed the OOD humans in proximity to the upper arm distribution.
4.6 Limitations and Failure Cases
Although PALM achieves good average-case performance, it works best with humans sampled near the training distribution.
If we pair the robot with an adversarial human, PALM is likely to fail as it lacks a fall-back safety policy.
The major limitation of PALM is the requirement of generating a human population.
While we provide one way to generate human populations based on weighted human-robot co-optimization, we lack ways to systematically generate diverse and realistic human motions.
One important direction for future work is to incorporate real user data to create training populations.
Improving the realism of the training human population is likely a crucial step to supporting transfer to real partners.
One future direction is extending to settings with one patient and one human caregiver.
While our framework still applies, this leads to new challenges including (1) learning a joint or separate latent space for human patient and caregiver, (2) modeling a population of human caregivers for training in simulation, and (3) modeling communication between the human caregiver and patient.
Conclusion
Generalization is an important task for assistive robotics, and in this paper, we formulate a problem setting that focuses on Out-Of-Distribution users.
To that end, we contribute a framework PALM for learning a robot policy that can quickly adapt to new partners at test time.
PALM assumes a distribution of training humans and constructs an embedding space for them by learning to predict partner actions.
We can further adapt this embedding at test time for new partners.
Experiments show that PALM outperforms state-of-the-art approaches.
We are excited by the potential of using PALM to enable robotic assistance in the future.
Acknowledgments
We would like to thank Ashish Kumar for discussions on the RMA method and Annie Xie for providing the implementation for LILI.
We would also like to thank Charlie C. Kemp for feedbacks and insights on assistive tasks.
This research was supported by the NSF National Robotics Initiative