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 𝜋 𝑅 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 (cid:104) 𝑆,𝛼, 𝐴 𝑅 , 𝐴 𝐻 , T , Ω 𝑅 , Ω 𝐻 ,𝑂, 𝑅 (cid:105) .
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
𝜋
𝐻 : Ω ∗ 𝐻 × 𝐴 𝐻
→
[ 0 , 1 ] be the human policy that maps from local histories of observations o
𝐻𝑡 = ( 𝑜 𝐻 1 ,...,𝑜 𝐻𝑡 ) over Ω 𝐻 to actions.
We deﬁne two distributions of human policies 𝜋 𝐻 ∈ 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 𝜋 𝐻 ∼ 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 𝜋 𝑅 : Ω ∗ 𝑅 → 𝐴 𝑅 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 𝜋 𝑅 = argmax 𝜋 𝐽 ( 𝜋,𝜋 𝐻 ) ,𝜋 𝐻 ∼ 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 handle 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 𝑧 ∼ E 𝜃 ( 𝑧 ; 𝜏 ) , where E 𝜃 encodes the trajectory 𝜏 so far and outputs latent vector 𝑧 . The robot uses this latent space to compute its actions 𝑎 𝑅𝑡 = 𝜋 𝑅 ( 𝑜 𝑅𝑡 ,𝑧 𝑡 ) . We train the base policy 𝜋 𝑅 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 𝑧 .
See Appendix E on training details. Ideally, we would like the latent space to capture suﬃcient information to diﬀerentiate the humans, similar to performing a “dimensionality reduction” on human policies 𝜋 𝐻 . We hereby introduce diﬀerent objectives for training 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, and in the assistive setting, we do not have access to humans reward function — 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 well 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.
We outline our method in Fig.
Given a history, 𝜏 1: 𝑡 = (cid:0) ( 𝑜 𝑅 1 ,𝑎 𝐻 1 ) ,... ( 𝑜 𝑅 𝑡 ,𝑎 𝐻 𝑡 ) (cid:1) , of robot observation and human action 2 pairs up to time 𝑡 , we embed this trajectory to a low-dimensional manifold.
The goal is to extract information about the human policy 𝜋 𝐻 .
Achieving this means that the vector 𝑧 is representative of the trajectory so far and indicative of the person’s future actions.
In other words, we can use 𝑧 to predict the person’s action at 𝑡 + 1.
We do this by training a decoder D 𝜙 parameterized by 𝜙 to predict the next action from the encoder’s output 𝑧 ∼ E 𝜃 ( 𝑧 ; 𝜏 1: 𝑡 ) .
[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 distribution.
This encourages nearby terms in the latent space to follow 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, 𝑧 ∗ of that user.
In other words, we would like to minimize ||E( 𝜏 ) − 𝑧 ∗ || 2 .
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, 𝜏 , and perform a few gradient steps to optimize both the encoder and decoder for Eq. (2): ( 𝜃,𝜙 ) → ( 𝜃,𝜙 ) − 𝛿 ∇ ( 𝜃,𝜙 ) L pred (E 𝜃 , D 𝜙 ,𝜏 ) .
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. We include an additional assistive task and a pilot user study in the appendix.
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 ( 𝑑 cos 𝛼 H ,𝑑 sin 𝛼 H ) , where 𝑑 is a ﬁxed value, and 𝛼 H ∈ [− 𝜋,𝜋 ] is known to the human, but not the robot.
The human agent is initialized randomly in the 2D plane with random hidden parameters 𝛼 H ∈ [− 𝜋,𝜋 ] , 𝑘 H ∈ [ 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 𝛼 H + 𝜋 2 , sin 𝛼 H + 𝜋 2 ) .
Only by recognizing the human in terms of 𝛼 H , 𝑘 H can the robot compensate the force, and successfully move the human to the hidden target. Each episode has 40 timesteps. The Scope of Generalization.
We deﬁne D train as 36 samples uniformly sampled from 𝛼 H ∈ [− 𝜋, 𝜋 2 ] , 𝑘 H ∈ [ 0 . 5 , 1 . 5 ] and D test as 12 samples uniformly sampled from 𝛼 H ∈ [ 𝜋 2 ,𝜋 ] , 𝑘 H ∈ [ 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 = 𝑐 p · || 𝜋 H ( 𝑠 𝑡 )
|| 2 where 𝑐 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 𝑐 p = 3 , 3 . 5 , 4, where larger penalties lead to the human agent exerting less eﬀort.
(2) 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 𝑐 p = 3.
In D 3 , we combine humans of 𝑐 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-Vanilla and RMA-Onehot , where we assign each training human a one-hot vector.
For RMA-Vanilla, we use a three-dimensional vector that includes the 𝑥, 𝑦 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 𝑐 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 𝑐 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 𝛼 H ∈ [− 𝜋,𝜋 ] 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 𝛼 H ,𝐾 𝐻 , the ideal embedding space looks like a ring with a small blob in the center.
The ring corresponds to the 2D projection of 𝛼 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 𝛼 H , it can learn a latent space that characterizes 𝛼 H .
Interestingly, varying the amount of regularization qualitatively aﬀects the shape of the latent space.
Setting the VAE regularization 𝑐 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 using Sec. for 200 epochs 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 𝑐 KL regularization is helpful for generalization.
Compare the latent space of PALM 𝑐 KL = 0 and PALM 𝑐 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
Supplementary Material
In this supplementary material, We ﬁrst provide details of a pilot user study in Sec.
A. We then present details and visualizations on how we generated D train , the training distribution of human agents in assistive itch scratching in Sec.
B. We then present the main algorithm pseudo-code for PALM in
Sec.
D.
Next we provide implementation details of diﬀerent methods in our main experiments in Sec.
E and the ablation study of PALM in Sec.
F and an evaluation of PALM with a second task, bed bathing assistance, in Sec.
G.
A User Study
We conduct a user study with 9 users from age 22 to 29.
We simplify the VR environment such that the user has four degrees of freedom: three on the shoulder and one on the elbow.
We do not explicitly track the elbow and use the controller rotation and orientation to track the arm motions.
We then convert the human pose deltas into actions in the environment.
Note that the real user can exceed the action limit in assistive gym ( env.action_space_human.high ).
This can be challenging for the robot because it has not seen such action magnitude during training.
We use a red marker to remind the users if they have exceeded the action limit.
For each user, we test with two itch spots, one on the upper arm and one on the lower arm.
We use randomized orderings of diﬀerent methods.
Each interaction is 150 simulation time steps and roughly six seconds of wall clock time.
Within each method, we let the user interact for a total of ﬁve trials.
We also ﬁnd that the robot can be initialized in poses that collide with the human at that moment.
We discard trials when such collisions are detected.
For training the robot controller, we synthesize humans of six diﬀerent activity levels ( 𝑐 p = 0 , 3 , 10 , 30 , 60 , 100), each with two random seeds to simulate users of diﬀerent activity levels.
We then sample 16 itch spots from each human, which results in a total of 192 training humans.
The training itch spots do not include the two spots that we test on.
Figure
For evaluation, we compute success count : the total number of time steps where the robot’s tool is within 10 centimeters of the itch location.
To account for diﬀerent user dexterity levels, we then normalize the result against the highest and the lowest success count on that itch spot for that user.
We show the results in Fig.
Note that PALM yields high performance in both itch spots and outperforms the RMA baselines.
We also ﬁnd the performance of the RNN baseline to be comparable with PALM for the ﬁrst task.
We hypothesize that this might be due to the real users being able to quickly adapt to the robot’s behaviour.
This causes the performance of all the methods to be more similar than shown in the simulation experiments.
B Generating Human Populations
To train our robot using sim2real, we would like to have a set of diverse environments.
However, unlike single-agent domain randomization where we can vary environment parameters such as friction, in assistive tasks the environment entails a changing user policy.
It is not obvious how to best generate a diverse population that captures user preferences, levels of disabilities, or movement characteristics.
Generating human motions that realistically capture the variation observed in physical human-robot interaction has remained an unsolved challenge in robotics.
Co-Optimization Prior works in robotic assistance [3, 25] have demonstrated that by optimizing for the same task objective, we can generate human and robot motions that coordinate towards the same goal, such as robot-assisted dressing.
Further more, we can leverage reward engineering [26, 27] to generate a diverse set of motions.
To generate diverse population in itch scratching task, we explore two sources of diversities: (1) we assign diﬀerent human action penalty 𝑐 p , where larger penalties lead to the human agent exerting less eﬀort.
In the simulation experiment we use 𝑐 p = 3 , 3 . 5 , 4 and in the user study we use 𝑐 p = 0 , 3 , 10 , 30 , 60 , 100.
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 as the ﬁ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).
Visualization Here we visualize trajectories from humans with diﬀerent action penalties in Fig.
Note that higher penalties result in the human taking more steady actions of smaller magnitude.
C Comparison with other VAE baselines for sequential data
Note that our method relies on embedding human trajectories as sequential data into a latent space.
Our implementation uses the ﬁnal hidden state of RNN as the input to variational autoencoder.
This is based on [28], which has been shown to be eﬀective in embedding and generating sentences.
Given that there are other diﬀerent generative models for sequential data, our framework can be easily combined with them.
In fact, wee believe coming up with a better model for human embedding is a future direction.
We hereby provide comparison with another diﬀerent generative sequential model [29].
Diﬀerent from [28] that uses only the ﬁnal hidden state, they construct a latent space for every intermediate step in the sequential model.
We keep all the experiment hyper-parameters the same, and concatenate the ﬁnal hidden state with observation as input to the robot policy.
We show the results in
Normalized Reward
Distribution
Our Method [28]
[Table]
D Algorithm
We present the main algorithm for PALM assume we have access to training and test distributions D train and D test .
Algorithm 1 Prediction-based Assistive Latent eMbedding Training
Randomly initialize base policy 𝜋 , encoder E 𝜙 parameterized by 𝜙 , decoder D 𝜃 parameterized by 𝜃 .
Empty replay buﬀer 𝐷 1 , window size 𝑤 for itr = 1 ,.., 𝑁 itr do for 𝑖 = 1 ,.., 𝑁 batch do
Sample 𝜋 H 𝑖 ∼ D train , optionally ﬁnd pre-trained expert robot policy
𝜋 E 𝑖 𝑜 0 ← env.reset()
Initialize history 𝐻 ← 𝜙 for 𝑡 = 0 ,...,𝑇 do
Get latest 𝑤 steps from 𝐻 : 𝐻 − 𝑤 ← 𝐻 [− 𝑤 : ] .
𝑧 𝑡 ← E 𝜙 ( 𝑜 𝑡 ,𝐻 − 𝑤 )
𝑎
H 𝑡 ← 𝜋 H 𝑖 ( 𝑜 0 )
𝑎
R 𝑡 ← 𝜋 ( 𝑜 𝑡 ,𝑧 𝑡 ) ,𝑎 e 𝑡 ← 𝜋 E 𝑖 ( 𝑜 𝑡 ) 𝑜 𝑡 + 1 ← env.step( 𝑎 H 𝑡 ,𝑎 R 𝑡 )
Store ( 𝑜 𝑡 ,𝑎 H 𝑡 ,𝑎 R 𝑡 ,𝑎 e 𝑡 ,𝐻 − 𝑤 ) in 𝐷end endfor 𝑗 = 1 ,..., 𝑁 opt do
Sample a batch of ( 𝑜 𝑡 ,𝑎 H 𝑡 ,𝑎 R 𝑡 ,𝐻 − 𝜏 ) from 𝐷
Compute L pred using Eq.
L PPO using [30] and L BC = (cid:205)
𝑡 || 𝑎 exp 𝑡 − 𝑎 𝑅𝑡
||
Optimize 𝜃,𝜙,𝜋 for L PALM =
𝜆 pred L pred + 𝜆 PPO L PPO + 𝜆 BC L BC endend
Sample 𝜋 H ∼ D test , Initialize history 𝐻 ← 𝜙 , Empty trajectory data 𝜏 for 𝑖 = 0 ,.., 𝑁 adapt do 𝑜 0 ← env.reset() for 𝑡 = 0 ,...,
𝑇 do
Get latest 𝑤 steps from 𝐻 : 𝐻 − 𝑤 ← 𝐻 [− 𝑤 : ] .
𝑧 𝑡 ← E 𝜙 ( 𝑜 𝑡 ,𝐻 − 𝑤 )
𝑎
R 𝑡 ← 𝜋 ( 𝑜 𝑡 ,𝑧 𝑡 ) 𝑜 𝑡 + 1 ← env.step( 𝑎 H 𝑡 ,𝑎 R 𝑡 )
Store ( 𝑜 𝑡 ,𝑎 H 𝑡 ,𝐻 − 𝑤 ) in 𝜏 end
Compute L pred using Eq.
(2) on 𝜏 , 𝜃 → 𝜃 − 𝛿 ∇ 𝜃 L pred (E 𝜃 ,𝜏 ) .
end
E Additional Training and Implementation Details
In our experiments in both the assistive reaching and assistive itch scratching, we use a recurrent network over a sliding window of 4-time steps, each of which is a concatenated vector of observation 𝑜 𝑡 , human action 𝑎 H 𝑡 − 𝑖 and robot actions 𝑎 R 𝑡 − 𝑖 .
For the current time step 𝑡 , we use zero vector for 𝑎
H 𝑡 and 𝑎 R 𝑡 .
We set the latent space dimension to be four, and use a recurrent network with six layers.
Our base policy network has four dimensions and hidden size of 100.
E.1 PPO and Behaviour Cloning
We need to train the base policy 𝜋 𝑅 and the latent space encoder E 𝜃 ( 𝑧 ; 𝜏 ) jointly because they are interdependent — 𝑧 is the input to 𝜋 𝑅 , and 𝜋 𝑅 decides the data distribution which leads to 𝑧 .
We simultaneously optimize the prediction loss in Eq.
(2) and the policy loss using PPO [30] algorithm.
See Appendix E for more training details.
To amend for the instability of training with a population of humans, we leverage Behaviour Cloning, where we use expert robot policies obtained via co-optimization (see Appendix B) to supervise 𝜋 𝑅 on on-policy data.
More speciﬁcally, we query the expert actions 𝑎 exp 𝑡 in a DAgger fashion[31] 3 during training, and optimize 𝑎 𝑅𝑡 to minimize deviation from it: L BC = (cid:205)
𝑡 || 𝑎 exp 𝑡 − 𝑎 𝑅𝑡
|| 2 .
The overall policy optimization loss include three terms: latent prediction loss, PPO loss and the Behaviour Cloning loss: L PALM =
L pred + L PPO + L BC .
E.2 Hyperparameters
PALM Training We use 𝜆 PPO = 0 .
1 ,𝜆 BC = 1 ,𝜆 pred = 0 .
1 in our experiments.
We ﬁnd that the behaviour cloning loss is essential for the Assistive Itch Scratching task.
Under this hyperparameter setting, we train for 200 iterations.
During each iteration, we collect 19,200 state-action transitions, which is evenly divided into 20 mini-batches.
Each mini-batch is fed to the base policy and encoder for 30 rounds to compute the loss and error for back-propagation.
We set the learning rate to be 0 . 00005.
For PALM prediction training, we use a three-layer decoder with hidden size 12 to predict the next human action from the hidden state from the encoder.
To implement the KL regularization, follow the standard VAE approach.
We use two linear networks to transform the encoder hidden state into 𝜇 and 𝜎 , which denote the mean and the standard deviation of the latent space.
We then compute approximate KL divergence to normal distribution on this latent distribution.
PALM Test Time Optimization At test time, we roll out the trained robot policy and collect data with the same user for 25 iterations, or 2,500 time steps.
This amounts to 150 seconds of wall clock time.
We then optimize for the prediction loss (including KL regularization term) using learning rate of 0.0001 for one to ﬁve steps, and use the one with the lowest loss.
We empirically ﬁnd the hyperparameters by doing the same process with humans from the training distribution, where we collect a mini training set and mini evaluation set both of 25 iterations.
We use the mini training set to ﬁnd the learning rate and use the evaluation set to ensure there is no over-ﬁtting.
RILI/LILI Training We follow a similar approach to PALM, except that we learn to predict the next state 𝑜 𝑡 + 1 and scalar reward.
RMA Training We follow the two-phase training procedure in [7].
Note that we ﬁnd it crucial in phase 2 to train the encoder with on-policy data, meaning that the regression data is collected by rolling out actions output by the “recurrent learner”, not the trained network from phase 1.
The phase 1 network is used simply for generating labels.
RNN/MLP Training .
For RNN, we directly feed the hidden state of the recurrent encoder to the policy network.
The architectural diﬀerence between RNN and PALM is that we do not concatenate the current observation 𝑜 𝑡 to the encoded output.
To ensure that the policy has at least the same capacity as PALM, we use a base policy with the same number of parameters as in PALM
[Table]
F Ablation Studies
[Table]
We include ablation studies of PALM in the main experiment in Sec. 4, where we study the eﬀect of test-time optimization, prediction loss, KL regularization and jointly training encoder E and policy 𝜋 .
In the assistive reaching experiment, we observe that test-time optimization, L pred , KL regularization, and joint training all contribute to the OOD performance.
In the assistive itch scratching experiment, test time optimization and L pred improve experiment results in all the settings.
Applying KL regularization provides some gain in the complex distribution D 3 , but does not lead to improvement in simpler distribution D 1 and D 2 .
G Additional Experiment: Bed Bathing Task
We further evaluate the performance of PALM in another assistive robotics task: robot-assisted bathing.
This task is a modiﬁed version of the bathing task introduced in Assistive Gym [3].
In this task, we have a human lying on a tilted bed, with a robot mounted on the nightstand.
The human can move their right arm, and there is an identiﬁed patch of skin to be cleaned.
Unlike the original bed bathing task, where the entire right arm is covered in target points to be cleaned, we instead initialize a ﬁxed size region of points to be cleaned at a uniform randomly selected location along the surface of the right forearm.
The patch spans 10 centimeter along and 150 degrees around the forearm.
Only the human knows the center position of the patch, and the robot must infer the location of the patch based on observations of the human motion.
The points along the body are cleaned whenever the robot initiates contact with the spot using its end-eﬀector and applies a positive normal force.
The task reward is based on how many points are cleaned.
We generate synthetic humans by co-optimizing humans and robots conditioned on the centroid position of the region to be cleaned.
During co-optimization, we adopt the formulation in [3] where both the human and the robot observe the centroid position along the human forearm.
This helps induce collaborative behaviour for the human and robot policies, where we then use the resulting human as the synthetic population.
During training, we blind of robot policy of the centroid position.
We use the co-optimized robot (observes the centroid position) as the oracle for querying expert actions for Behaviour Cloning.
We sample 18 humans from the training distribution, and save 6 humans from the held-out distribution as the out-of-distribution evaluation.
We use the same hyperparameters as the itch scratching experiment.
As we see in the results in Fig.
11, PALM achieves better out-of-distribution results compared to the baseline methods.
We also observe that this performance gap is smaller than the itch scratching task.
We believe this is due to the nature of the bed bathing task, where a robot controller that maintains contact with the human’s forearm can be suﬃcient for solving the task if the human policy learns to move and rotate their forearm accordingly to help the robot.