DARA: Dynamics-Aware Reward Augmentation in Offline Reinforcement Learning

1 INTRODUCTION Offline reinforcement learning ( RL ) ( Levine et al. , 2020 ; Lange et al. , 2012 ) , the task of learning from the previously collected dataset , holds the promise of acquiring policies without any costly active interaction required in the standard online RL paradigm . However , we note that although the active trail-and-error ( online exploration ) is eliminated , the performance of offline RL method heavily relies on the amount of offline data that is used for training . As shown in Figure 1 , the performance deteriorates dramatically as the amount of offline data decreases . A natural question therefore arises : can we reduce the amount of the ( target ) offline data without significantly affecting the final performance for the target task ? Bringing the idea from the transfer learning ( Pan & Yang , 2010 ) , we assume that we have access to another ( source ) offline dataset , hoping that we can leverage this dataset to compensate for the performance degradation caused by the reduced ( target ) offline dataset . In the offline setting , previous work ( Siegel et al. , 2020 ; Chebotar et al. , 2021 ) has characterized the reward ( goal ) difference between the source and target , relying on the ” conflicting ” or multi-goal offline dataset ( Fu et al. , 2020 ) , while we focus on the relatively unexplored transition dynamics difference between the source dataset and the target environment . Meanwhile , we believe that this dynamics shift is not arbitrary in reality : in healthcare treatment , offline data for a particular patient is often limited , whereas we can obtain diagnostic data from other patients with the same case ( same reward/goal ) and there often exist individual differences between patients ( source dataset with different transition dynamics ) . Careful treatment with respect to the individual differences is thus a crucial requirement . Given source offline data , the main challenge is to cope with the transition dynamics difference , i.e. , strictly tracking the state-action supported by the source offline data can not guarantee that the same transition ( state-action-next-state ) can be achieved in the target environment . However , in the offline setting , such dynamics shift is not explicitly characterized by the previous offline RL methods , where they typically attribute the difficulty of learning from offline data to the state-action distribution shift ( Chen & Jiang , 2019 ; Liu et al. , 2018 ) . The corresponding algorithms ( Fujimoto et al. , 2019 ; Abdolmaleki et al. , 2018 ; Yu et al. , 2020 ) that model the support of state-action distribution induced by the learned policy , will inevitably suffer from the transfer problem where dynamics shift happens . Our approach is motivated by the well established connection between reward modification and dynamics adaptation ( Kumar et al. , 2020b ; Eysenbach & Levine , 2019 ; Eysenbach et al. , 2021 ) , which indicates that , by modifying rewards , one can train a policy in one environment and make the learned policy to be suitable for another environment ( with different dynamics ) . Thus , we propose to exploit the joint distribution of state-action-next-state : besides characterizing the state-action distribution shift , we additionally identify the dynamics ( i.e. , conditional distribution of next-state given state-action pair ) shift and penalize the agent with a dynamics-aware reward modification . This modification aims to discourage the learning from these offline transitions that are likely in source but are unlikely in target . Unlike the concurrent work ( Ball et al. , 2021 ; Mitchell et al. , 2021 ) paying attention to the offline domain generalization , we explicitly focus on the offline domain ( dynamics ) adaptation . Our principal contribution in this work is the characterization of the dynamics shift in offline RL and the derivation of dynamics-aware reward augmentation ( DARA ) framework built on prior modelfree and model-based formulations . DARA is simple and general , can accommodate various offline RL methods , and can be implemented in just a few lines of code on top of dataloader at training . In our offline dynamics adaptation setting , we also release a dataset , including the Gym-MuJoCo tasks ( Walker2d , Hopper and HalfCheetah ) , with dynamics ( mass , joint ) shift compared to D4RL , and a 12-DoF quadruped robot in both simulator and real-world . With only modest amounts of target offline data , we show that DARA-based offline methods can acquire an adaptive policy for the target tasks and achieve better performance compared to baselines in both simulated and real-world tasks . 2 RELATED WORK . Offline RL describes the setting in which a learner has access to only a fixed dataset of experience , while no interactive data collection is allowed during policy learning ( Levine et al. , 2020 ) . Prior work commonly assumes that the offline experience is collected by some behavior policies on the same environment that the learned policy be deployed on . Thus , the main difficulty of such offline setting is the state-action distribution shift ( Fujimoto et al. , 2019 ; Liu et al. , 2018 ) . Algorithms address this issue by following the two main directions : the model-free and model-based offline RL . Model-free methods for such setting typically fall under three categories : 1 ) Typical methods mitigate this problem by explicitly ( Fujimoto et al. , 2019 ; Kumar et al. , 2019 ; Wu et al. , 2019 ) or implicitly ( Siegel et al. , 2020 ; Peng et al. , 2019 ; Abdolmaleki et al. , 2018 ) constraining the learned policy away from OOD state-action pairs . 2 ) Conservative estimation based methods learn pessimistic value functions to prevent the overestimation ( Kumar et al. , 2020a ; Xu et al. , 2021 ) . 3 ) Importance sampling based methods directly estimate the state-marginal importance ratio and obtain an unbiased value estimation ( Zhang et al. , 2020 ; Nachum & Dai , 2020 ; Nachum et al. , 2019b ) . Model-based methods typically eliminate the state-action distribution shift by incorporating a reward penalty , which relies on the uncertainty quantification of the learned dynamics ( Kidambi et al. , 2020 ; Yu et al. , 2020 ) . To remove this uncertainty estimation , Yu et al . ( 2021 ) learns conservative critic function by penalizing the values of the generated state-action pairs that are not in the offline dataset . These methods , however , define their objective based on the state-action distribution shift , and ignore the potential dynamics shift between the fixed offline data and the target MDP . In contrast , we account for dynamics ( state-action-next-state ) shift and explicitly propose the dynamics aware reward augmentation . A counterpart , close to our work , is off-dynamics RL ( Eysenbach et al. , 2021 ) , where they set up dynamics shift in the interactive environment while we focus on the offline setting . 3 PRELIMINARIES . We study RL in the framework of Markov decision processes ( MDPs ) specified by the tuple M : = ( S , A , r , T , ρ0 , γ ) , where S and A denote the state and action spaces , r ( s , a ) ∈ [ −Rmax , Rmax ] is the reward function , T ( s′|s , a ) is the transition dynamics , ρ0 ( s ) is the initial state distribution , and γ is the discount factor . The goal in RL is to optimize a policy π ( a|s ) that maximizes the expected discounted return ηM ( π ) : = Eτ∼pπM ( τ ) [ ∑∞ t=0 γ tr ( st , at ) ] , where τ : = ( s0 , a0 , s1 , a1 , ... ) . We also define Q-values Q ( s , a ) : = Eτ∼pπM ( τ ) [ ∑∞ t=0 γ tr ( st , at ) |s0 = s , a0 = a ] , V-values V ( s ) : = Ea∼π ( a|s ) [ Q ( s , a ) ] , and the ( unnormalized ) state visitation distribution dπM ( s ) : = ∑∞ t=0 γ tP ( s|π , M , t ) , where P ( s|π , M , t ) denotes the probability of reach state s at time t by running π in M . In the offline RL problem , we are provided with a static dataset D : = { ( s , a , r , s′ ) } , which consists of transition tuples from trajectories collected by running one or more behavioral policies , denoted by πb , on MDP M . With a slight abuse of notation , we write D = { ( s , a , r , s′ ) ∼ dD ( s ) πb ( a|s ) r ( s , a ) T ( s′|s , a ) } , where the dD ( s ) denotes state-marginal distribution in D. In the offline setting , the goal is typically to learn the best possible policy using the fixed offline dataset . Model-free RL algorithms based on dynamic programming typically perform policy iteration to find the optimal policy . Such methods iteratively conduct 1 ) policy improvement with GMQ : = arg maxπ Es∼dπM ( s ) , a∼π ( a|s ) [ Q ( s , a ) ] and 2 ) policy evaluation by iterating the Bellman equation Q ( s , a ) = BπMQ ( s , a ) : = r ( s , a ) + γEs′∼T ( s′|s , a ) , a′∼π ( a′|s′ ) [ Q ( s′ , a′ ) ] over dπM ( s ) π ( a|s ) . Given off-policy data D , we resort to 1 ) improvement with GDQ : = arg maxπ Es∼dD ( s ) , a∼π ( a|s ) [ Q ( s , a ) ] and 2 ) evaluation by iteratingQ ( s , a ) = BπDQ ( s , a ) : = r ( s , a ) +γEs′∼T ( s′|s , a ) , a′∼π ( a′|s′ ) [ Q ( s′ , a′ ) ] over all ( s , a ) in D. Specifically , given any initial Q0 , it iterates1 Policy improvement : πk+1 = GDQk , Policy evaluation : Qk+1 = Bπ k+1 D Q k. ( 1 ) Model-free offline RL based on the above iteration suffers from the state-action distribution shift , i.e. , policy evaluation BπkD Qk−1 may encounter unfamiliar state action regime that is not covered by the fixed offline dataset D , causing erroneous estimation of Qk . Policy improvement GDQk further exaggerates such error , biasing policy πk+1 towards out-of-distribution ( OOD ) actions with erroneously high Q-values . To address this distribution shift , prior works 1 ) explicitly constrain policy to be close to the behavior policy ( Fujimoto et al. , 2019 ; Kumar et al. , 2019 ; Wu et al. , 2019 ; Ghasemipour et al. , 2021 ) , introducing penalty αD ( π ( a|s ) , πb ( a|s ) ) into GD or BπD in Equation 1 : GDQ = arg max π Es∼dD ( s ) , a∼π ( a|s ) [ Q ( s , a ) − αD ( π ( a|s ) , πb ( a|s ) ) ] , BπDQ ( s , a ) = r ( s , a ) + γEs′∼T ( s′|s , a ) , a′∼π ( a′|s′ ) [ Q ( s′ , a′ ) − αD ( π ( a′|s′ ) , πb ( a′|s′ ) ) ] , ( 2 ) where D is a divergence function between distributions over actions ( e.g. , MMD or KL divergence ) , or 2 ) train pessimistic value functions ( Kumar et al. , 2020a ; Yu et al. , 2021 ; Xu et al. , 2021 ) , penalizing Q-values at states in the offline dataset D for actions generated by the current policy π : Q = arg min Q Es∼dD ( s ) , a∼π ( a|s ) [ Q ( s , a ) ] , s.t . Q = B π DQ . ( 3 ) Model-based RL algorithms iteratively 1 ) model the transition dynamics T ( s′|s , a ) , using the data collected in M : maxT̂ Es , a , s′∼dπM ( s ) π ( a|s ) T ( s′|s , a ) [ log T̂ ( s ′|s , a ) ] , and 2 ) infer a policy π from the modeled M̂ = ( S , A , r , T̂ , ρ0 , γ ) , where we assume that r and ρ0 are known , maximizing ηM̂ ( π ) with a planner or the Dyna-style algorithms ( Sutton , 1990 ) . In this paper , we focus on the latter . Model-based offline RL algorithms similarly suffer from OOD state-action ( Kidambi et al. , 2020 ; Cang et al. , 2021 ) if we directly apply policy iteration over T̂ : = maxT̂ Es , a , s′∼D [ log T̂ ( s ′|s , a ) ] . Like the conservative estimation approach described in Equation 3 , recent conservative model-based offline RL methods provide the policy with a penalty for visiting states under the estimated T̂ where T̂ is likely to be incorrect . Taking u ( s , a ) as the oracle uncertainty ( Yu et al. , 2020 ) that provides a consistent estimate of the accuracy of model T̂ at ( s , a ) , we can modify the reward function to obtain a conservative MDP : M̂c = ( S , A , r − αu , T̂ , ρ0 , γ ) , then learn a policy π by maximizing ηM̂c ( π ) .

Summary: This paper presents a new offline RL algorithm that aims to adapt the policy to a target domain with insufficient data with the help of abundant data in a source domain. The authors propose their method, DARA, which applies a reward penalty that correct the dynamics shift between the source and target domains. The authors combine DARA with previous model-based and model-free offline RL methods and show that DARA is able to achieve performance improvement over prior offline RL methods without the reward augmentation in both simulated and real-world tasks.