Unsupervised Domain Adaptation with Dynamics-Aware Rewards in Reinforcement Learning

1 Introduction . Recently , the machine learning community has devoted attention to unsupervised reinforcement learning ( RL ) to acquire useful skills , ie , the problem of automatic discovery of a goal-conditioned policy and its corresponding goal space [ 8 ] . As shown in Figure 1 ( left ) , the standard training procedure of learning skills in an unsupervised way follows : ( 1 ) representing goals , consisting of automatically generating the goal distribution p ( g ) and the corresponding goal-achievement reward function rg ; ( 2 ) learning the goal-conditioned policy πθ with the acquired p ( g ) and rg . Leveraging fully autonomous interaction with the environment , the agent sets up goals , builds the goal-achievement reward function , and extrapolates the goal-conditioned policy in parallel by adopting off-the-shelf RL methods [ 40 , 19 ] . While we can obtain skills without any prior goal representations ( p ( g ) and rg ) in an unsupervised way , a major drawback of this approach is that it requires a large amount of rollout steps to represent goals and learn the policy itself , together . This procedure is often impractical in some target environments ( eg , the robot in real world ) , where online interactions are time-consuming and potentially expensive . That said , there often exist environments that resemble in structure ( dynamics ) yet provide more accessible rollouts ( eg , unlimited in simulators ) . For problems with such source environments available , training the policy in a source environment significantly reduces the cost associated with interaction in the target environment . Critically , we can train a policy in one environment and deploy it in another by utilizing their structural similarity and the excess of interaction . Considering the navigation in a room , we can learn arbitrary skills through the active exploration in a source simulated ∗Work was done at Westlake University . † Corresponding author . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . room ( with different layout or friction ) before the deployment in the target room . However , it is reasonable to suspect that the learned skills overfit the training environment , the dynamics of which , dictating the goal distribution and reward function , implicitly shape goal representation and guide policy acquisition . Such deployment would then make learned skills struggle to adapt to new , unseen environments and produce a large drop in performance in target due to the dynamics shifts , as shown in Figure 2 ( top ) . In this paper , we overcome the limitations ( of limited rollout in target and dynamics shifts ) associated with the ( source , target ) environments pair through unsupervised domain adaptation . In practice , while performing a full unsupervised RL method in target that represents goals and captures all of them for learning the entire goal-conditioned policy ( Figure 1 left ) can be extremely challenging with the limited rollout steps , learning a model for only ( partially ) representing goals is much easier . This gives rise to learning the policy in source and taking the limited rollouts in target into account only for identifying the goal representations , which further shape the policy . As shown in Figure 1 ( right ) , we represent goals in both environments while optimizing the policy only in the source environment , alleviating the excessive need for rollout steps in the target environment . Furthermore , we introduce a KL regularization to address the challenge of dynamics shifts . This objective allows us to incorporate a reward modification into the goal-achievement reward function in the standard unsupervised RL , aligning the trajectory induced in the target environment against that induced in the source by the same policy . Importantly , it enables useful inductive biases towards the target dynamics : it allows the agent to specifically pursue skills that are competent in the target dynamics , and penalizes the agent for exploration in the source where the dynamics significantly differ . As shown in Figure 2 ( bottom ) , the difference in dynamics ( a wall in the target while no wall in the source ) will pose a penalty when the agent attempts to go through an area in the source wherein the target stands a wall . Thus , skills learned in source with such modification are adaptive to the target . We name our method unsupervised domain adaptation with dynamics-aware rewards ( DARS ) , suggesting that source and target dynamics both shape rg : ( 1 ) we employ a latent-conditioned probing policy in the source to represent goals [ 31 ] , making the goal-achievement reward source-oriented , and ( 2 ) we adopt two classifiers [ 11 ] to provide reward modification derived from the KL regularization . This means that the repertoires of skills are well shaped by the dynamics of both the source and target . Formally , we further analyze the conditions under which our DARS produces a near-optimal goal-conditioned policy for the target environment . Empirically , we demonstrate that our objective can obtain dynamics-aware rewards , enabling the goal-conditioned policy learned in a source to perform well in the target environment in various settings ( stable and unstable settings , and sim2real ) . 2 Preliminaries . Multi-goal Reinforcement Learning : We formalize the multi-goal reinforcement learning ( RL ) as a goal-conditioned Markov Decision Process ( MDP ) defined by the tupleMG = { S , A , P , RG , γ , ρ0 } , where S denotes the state space and A denotes the action space . P : S × A × S → R≥0 is the transition probability density . RG , { G , rg , p ( g ) } , where G denotes the space of goals , rg denotes the corresponding goal-achievement reward function rg : G× S ×A× S → R , and p ( g ) denotes the given goal distribution . γ is the discount factor and ρ0 is the initial state distribution . Given a g ∼ p ( g ) , the γ-discounted return R ( g , τ ) of a goal-oriented trajectory τ = ( s0 , a0 , s1 , . . . , sT ) is ΣT−1t=0 γ trg ( st , at , st+1 ) . Building on the universal value function approximators ( UVFA , Schaul et al . [ 38 ] ) , the standard multi-goal RL seeks to learn a unique goal-conditioned policy πθ : A×S×G→ R to maximize the objective EP , ρ0 , πθ , p ( g ) [ R ( g , τ ) ] , where θ denotes the parameter of the policy . Unsupervised Reinforcement Learning : In unsupervised RL , the agent is set in an open-ended environment without any pre-defined goals or related reward functions . The agent aims to acquire a repertoire of skills . Following Colas et al . [ 8 ] , we define skills as the association of goals and the goalconditioned policy to reach them . The unsupervised skill acquisition problem can now be modeled by a goal-free MDPM = { S , A , P , γ , ρ0 } that only characterizes the agent , its environment and their possible interactions . As shown in Figure 1 ( left ) , the agent needs to autonomously interact with the environment and ( 1 ) learn goal representations ( eg , discovering the goal distribution p ( g ) and learning the corresponding reward rg ) , and ( 2 ) learn the goal-conditioned policy πθ as in multi-goal RL . Here we define a universal ( information theoretic ) objective for learning the goal-conditioned policy πθ in unsupervised RL , maximizing the mutual information IP , ρ0 , πθ ( g ; τ ) between the goal g and the trajectory τ induced by policy πθ running in the environmentM ( with P and ρ0 ) , max IP , ρ0 , πθ ( g ; τ ) = H ( g ) −H ( g|τ ) = H ( g ) + EP , ρ0 , πθ , p ( g ) [ log p ( g|τ ) ] . ( 1 ) For representing goals , the specific manifold of the goal space could be a set of latent variables ( eg , one-hot indicators ) or perceptually-specific goals ( eg , the joint torques of ant ) . In the absence of any prior knowledge about p ( g ) , the maximum ofH ( g ) will be achieved by fixing the distribution p ( g ) to be uniform over all g ∈ G. The second term EP , ρ0 , πθ , p ( g ) [ log p ( g|τ ) ] in Equation 1 is analogous to the objective in the standard multi-goal RL , where the return R ( g , τ ) can be seen as the embodiment of log p ( g|τ ) . The objective specifically for learning rg in p ( g|τ ) is normally optimized by lens of the generative loss [ 33 ] or the contrastive loss [ 42 ] . With the learned goal distribution p ( g ) and reward rg , it is straightforward to learn the goal-conditioned policy πθ using standard RL algorithms [ 40 , 19 ] . In general , optimizations iteratively alternate for representing goals ( including both goal-distribution p ( g ) and reward function rg ) and learning the goal-conditioned policy πθ , as shown in Figure 1 ( left ) . 3 Unsupervised Domain Adaptation with Dynamics-Aware Rewards . 3.1 Problem Formulation . Our work addresses domain adaptation in unsupervised RL , raising expectations that an agent trained without prior goal representations ( p ( g ) and rg ) in one environment can perform purposeful tasks in another . Following Wulfmeier et al . [ 54 ] , we also focus on the domain adaptation of the dynamics , as opposed to states . In this work , we consider two environments characterized by MDPs MS ( the source environment ) andMT ( the target environment ) , the dynamics of which are PS and PT respectively . Both MDPs share the same state and action spaces S , A , discount factor γ and initial state distribution ρ0 , while differing in the transition distributions PS , PT . Since the agent does not directly receiveRG from either environment , we adopt the information theoretic IP , ρ0 , πθ ( g ; τ ) to acquire skills , equivalently learning a goal-conditioned policy πθ that achieves distinguishable trajectory by maximizing this objective . For brevity , we now omit the ρ0 term discussed in Section 2 . In our setup , agents can freely interact with the sourceMS . However , it has limited access to rollouts in the targetMT with which are insufficient to train a policy . To ensure that all potential trajectories in the targetMT can be attempted in the source environment , we make the following assumption : Assumption 1 . There is no transition that is possible in the target environmentMT but impossible in the source environmentMS : PT ( st+1|st , at ) > 0 =⇒ PS ( st+1|st , at ) > 0 . 3.2 Domain Adaptation in Unsupervised RL . We aim to acquire skills trained in the source environmentMS , which can be deployed in the target environmentMT . To facilitate the unsupervised learning of skills for the target environmentMT ( with transition dynamics PT ) , we maximize the mutual information between the goal g and the trajectory τ induced by the goal-conditioned policy πθ over dynamics PT , as shown in Figure 3 ( a ) : IPT , πθ ( g ; τ ) . ( 2 ) However , since interaction with the target environment MT is restricted , acquiring the goalconditioned policy πθ by optimizing the mutual information above is intractable . We instead maximize the mutual information in the source environment IPS , πθ ( g ; τ ) modified by a KL divergence of trajectories induced by the goal-conditioned policy πθ in both environments ( Figure 3 b ) : IPS , πθ ( g ; τ ) − βDKL ( pPS , πθ ( g , τ ) ‖pPT , πθ ( g , τ ) ) , ( 3 ) where β > 0 is the regularization coefficient , pPS , πθ ( g , τ ) and pPT , πθ ( g , τ ) denote the joint distributions of the goal g and the trajectory τ induced by policy πθ in sourceMS and targetMT respectively . Intuitively , maximizing the mutual information term rewards distinguishable pairs of trajectories and goals , while minimizing the KL divergence term penalizes producing a trajectory that can not be followed in the target environment . In other words , the KL term aligns the probability distributions of the mutual-information-maximizing trajectories under the two environment dynamics PS and PT . This indicates that the dynamics of both environments ( PS and PT ) shape the goal-conditioned policy πθ ( even though trained in the source PS ) , allowing πθ to adapt to the shifts in dynamics . Building on the KL regularized objective in Equation 3 , we introduce how to effectively represent goals : generating the goal distribution and acquiring the ( partial ) reward function . Here we assume the difference between environments in their dynamics negligibly affects the goal distribution2 . Therefore , we follow GPIM [ 31 ] and train a latent-conditioned probing policy πµ . The probing policy πµ explores the source environment and represents goals for the source to train the goal-conditioned policy πθ with . Specifically , the probing policy πµ is conditioned on a latent variable ω ∼ p ( ω ) 3 and aims to generate diverse trajectories that are further relabeled as goals for the goal-conditioned πθ . Such goals can take the form of the latent variable ω itself ( Figure 3 c1 ) or the final state of a trajectory ( Figure 3 c2 ) . We jointly optimize the previous objective in Equation 3 with the mutual information between ω and the trajectory τ̃ induced by πµ in source , and arrive at the following overall objective : maxJ ( µ , θ ) , IPS , πµ ( ω ; τ̃ ) + IPS , πθ ( g ; τ ) − βDKL ( pPS , πθ ( g , τ ) ‖pPT , πθ ( g , τ ) ) , ( 4 ) where the context between p ( g ) and p ( ω ) are specified by the graphic model in Figure 3 ( c1 or c2 ) . Note that this objective ( Equation 4 ) explicitly decouples the goal representing ( with πµ ) and the policy learning ( wrt πθ ) , providing a foundation for the theoretical guarantee in Section 3.4 .

Summary: The authors address the problem of unsupervised Domain Adaptation in Reinforcement Learning. Combining ideas from previous work, they devise a method able to utilize an easy to sample source environment (e.g., a simulator) in order to learn a set of skills without supervision, while at the same time guarantee that these skills will be suited to the target environment (e.g., a real robotic system). To achieve this, they modify the reward structure of the typical Unsupervised Reinforcement Learning setup to enable the agent to build skills that are compatible with the target environment dynamics, while penalizing exploration in regions of the source environment where the dynamics differ. The proposed approach is benchmarked against state-of-the-art algorithms in a wide variety of environments with increased complexity and manages to outperform previous baselines both in simulation as well as in three walking tasks on a real quadruped robot.