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: The given work proposes the addition of a dynamics aware framework for Offline Reinforcement Learning. It motivates that  offline reinforcement learning suffer from a dynamics shift in SARSA transitions, which need to be effectively handled. They propose modifying reward in source offline dataset, penalizing it with the dynamics aware reward transition, setting the lower bound. The reward augmentation is done with explicit policy constraints as well as implicit policy constraints.They evaluate this approach on various robotic manipulation task from the Gym-MujoCo suite, evaluating and showcasing perforomance on four offline RL algorithms.