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: This paper presents an approach for unsupervised learning of a goal-conditioned policy that is robust to changes in dynamics and can be transferred to a target domain that is different from the source domain in terms of the underlying dynamics. There are several components. First, a latent-conditioned probing policy is used to generate goals (either in the source or target domain). This probing policy is trained to generate diverse goals by maximising the information gain between the goals and the resulting trajectories generated by the policy, i.e. we want diverse goals that are predictable given the resulting trajectories. This “predictability” is approximated using a discriminator that scores the likelihood of the goal latent given the resulting trajectory; this is also trained alongside the goal-generating policy based on rollouts in the source domain.   Next, the probing policy is used to generate goals in the source and target domains (these are either the latents themselves or the final states of trajectories from the probing policy), and the goal-conditioned policy (which is conditioned on these goals) is executed in both these domains to generate rollouts. A key assumption here is that it is costly to execute the policy in the target domain (e.g. real world), so the approach tries to minimise this by training the probing policy, discriminator and to a major extent, the policy, primarily on the source domain. The policy is again trained to maximise the log-likelihood of the discriminator, thereby leading to predictable trajectories in the source domain. To ensure that the policy does not overfit to the source domain, and consequently fail on the target domain, an additional KL regularisation term is added to the policy training objective. This regularisation penalises trajectories that provide inconsistent results between the source and target domains, as measured by the KL divergence of the resulting transition distributions when executing the policy for the same goal in both these domains. This effectively ensures that the policy learns trajectories that are executable consistently in both domains, thereby ensuring success even in the presence of differences in dynamics between the two domains. Note that two key assumptions that are necessary for this to work (and the theoretical results) is that there is no transition in the target domain that is not possible in the source domain, and that the goal distribution in the source is also valid in the target (a relaxation of the latter is provided in the supplementary material). This regularisation term is approximated by a pair of classifiers that are trained on data from both domains; these are also trained alongside the policy. The final algorithm alternates training the probing policy & discriminator that parameterise the goal distribution and rewards against training the goal-conditioned policy and the classifiers for approximating the dynamics difference KL.  The approach is tested on several simulated tasks of varying complexity. First, the approach is tested on a 2D point mass with pairs of source and target tasks where the arena has walls of different configurations. Next, the approach is evaluated on continuous control tasks from the OpenAI Gym with the ant, half-cheetah and humanoid embodiments. The source tasks are the standard gym tasks and the target tasks have certain joints not working or additional perturbations in the form of obstacles. Lastly, the approach is also tested in a sim2real setup with a quadruped; the source task is in simulation while the target task is on the real robot. The resulting policy and goal distribution show good performance on all these tasks; the skills and goal distribution learned by the method is able to successfully transfer across changes in wall configuration, joint issues, stability problems and also to real world robots. Several visualisations of different parts of the algorithm (e.g. discriminator, skills) etc. are provided along with comparisons to baselines, namely variants of DARC, SMiRL, DIAYN and GPIM. The supplement also provides additional analysis and proofs of the theoretical guarantees of the near-optimality of the proposed method under the pertinent assumptions.