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 paper handles the problem of domain adaptation in RL, specifically the setting where rewards and prior goal representations are not available. They do so by combining a skill learning method (in the source environment) with a transfer learning method that aims to match the joint $(goal, \tau)$ distribution arcoss environments. By doing so they show that skills can both be learned in the source domain, and they transfer more easily to the target domain.