A Policy-Guided Imitation Approach for Offline Reinforcement Learning
Abstract
Offline reinforcement learning (RL) methods can generally be categorized into two types: RL-based and imitation-based.
RL-based methods could in principle enjoy out-of-distribution generalization but suffer from erroneous off-policy evaluation.
Imitation-based methods avoid off-policy evaluation but are too conservative to surpass the dataset.
In this study, we propose an alternative approach, inheriting the training stability of imitation-style methods while still allowing logical out-ofdistribution generalization.
We decompose the conventional reward-maximizing policy in offline RL into a guide-policy and an execute-policy.
During training, the guide-policy and execute-policy are learned using only data from the dataset, in a supervised and decoupled manner.
During evaluation, the guide-policy guides the execute-policy by telling where it should go so that the reward can be maximized.
By doing so, our algorithm allows state-compositionality from the dataset, rather than action-compositionality conducted in prior imitation-style methods.
We dumb this new approach Policy-guided Offline RL ( POR ).
POR demonstrates the stateof-the-art performance on D4RL, a standard benchmark for offline RL.
We also highlight the benefits of POR in terms of improving with supplementary suboptimal data and easily adapting to new tasks by only changing the guide-policy.
Code is available at https://github.com/ryanxhr/POR .
Introduction
Offline RL, also known as batch RL, allows learning policies from previously collected data, without online interactions [30, 32].
It is a promising area for bringing RL into real-world domains, such as robotics [22], healthcare [48] and industrial control [63].
In such scenarios, arbitrary exploration with untrained policies is costly or dangerous but sufficient prior data is available.
While most off-policy
RL algorithms are applicable in the offline setting by filling the replay buffer with offline data, improving the policy beyond the level of the behavior policy often entails querying the value function (i.e., Q function) about values of actions that were not seen in the dataset.
Those out-of-distribution actions can be deemed as adversarial examples of the Q function [29], which cause extrapolation error of the Q-function.
The error will be accumulated by the deadly triad issue [51], propagate across the state-action space through the iterative dynamic programming procedure [5], and can not be eliminated without requiring a growing batch of online samples [32].
To alleviate this issue, prior model-free offline RL methods typically add a behavior regularization term to the policy improvement step, to limit how far it deviates from the behavior policy.
This can be achieved explicitly by calculating some divergence metrics [27, 55, 39, 14], or implicitly by regularizing the learned value functions to assign low values to out-of-distribution actions [29, 25, 2, 3].
Nevertheless, this imposes a trade-off between accurate value estimation (more behavior regularization) and maximum policy improvement (less behavior regularization).
To avoid this problem, recently a branch of methods bypass querying the values of unseen actions by performing some kind of imitation learning on the dataset * .
This can be achieved by filtering trajectories based on their return [41, 9], or reweighting transitions based on how advantageous they could be under the behavior policy [5, 26], or just be conditioned on some variables without any dataset reweighting [8, 11].
Although imitation-style methods enjoy a stable training process and are able to effectively perform multi-step dynamic programming by assigning the proper reweighting weight or conditioned variable [26, 11], they only allow action-compositionality from the dataset, lose the ability to surpass the dataset by out-of-distribution generalization, which only appears in RL-based methods.
In this work, we propose an alternative approach.
We aim at inheriting the training stability of imitation-style methods while still allowing logical out-of-distribution generalization.
To do so, we propose Policy-guided Offline RL ( POR ), an algorithm that employs two policies to solve tasks: a guide-policy and an execute-policy.
The job of the guide-policy is to learn the optimal next state given the current state, and the job of the execute-policy is to learn how different actions can produce different next states, given the current state.
By this manner, we decompose the original task of
RL (i.e., reward-maximizing) into two distinct yet complementary tasks, and this decomposition makes each task much easier to be solved.
During evaluation, the guide-policy serves as a guide for the execute-policy by telling it where to go so that the reward can be maximized.
By doing so, our algorithm allows state-compositionality from the dataset, rather than action-compositionality conducted in prior work, which owns logical out-of-distribution generalization.
Our method is easy to implement by only adding the learning of a guide-policy, the training stage of the guide-policy and execute -policy are decoupled, using in-sample learning from the dataset.
We test
POR in widely-used D4RL offline RL benchmarks and demonstrates the state-of-the-art performance, especially on low-quality datasets that require out-of-distribution generalization to achieve a high score.
We also show that by decoupling the learning of guide-policy and execute-policy, we can enhance the guide-policy with supplementary suboptimal data, or re-learn the guide-policy to adapt to new tasks with different reward functions, without changing the execute-policy.
Related Work RL-based offline approach
A large portion of offline RL methods is RL-based.
RL-based methods typically augment existing off-policy methods (e.g., Q-learning or actor-critic) with a behavior regularization term.
The primary ingredient of this class of methods is to propose various policy regularizers to ensure that the learned policy does not stray too far from the behavior policy.
These regularizers can appear explicitly as divergence penalties [55, 27, 14], implicitly through weighted behavior cloning [53, 41, 39], or more directly through careful parameterization of the policy [15, 66].
Another way to apply behavior regularizers is via modification of the critic learning objective to incorporate some form of regularization to encourage staying near the behavioral distribution and being pessimistic about unknown state-action pairs [38, 29, 25, 59].
There are also several works incorporating behavior regularization through the use of uncertainty.
The uncertainty quantification can be done via Monte Carlo dropout [56] or explicit ensemble models [3, 62, 23, 64].
Note that the behavior regularization weight is crucial in RL-based methods.
With a small weight, RL-based methods could in principle enjoy out-of-distribution generalization since they perform true dynamic programming.
However, a small weight will make erroneous off-policy evaluation due to distribution shift and the policy is extremely vulnerable to the value function misestimation.
Imitation-based offline approach
Another line of methods, on the contrary, performs some kind of imitation learning on the dataset, without the need to do off-policy evaluation.
When the dataset is good enough or contains high-performing trajectories, we can simply clone the actions observed in the dataset [42], or perform some kind of filtering or conditioning to extract useful transitions.
For instance, recent work filters trajectories based on their return [9, 41], or directly filters individual transitions based on how advantageous these could be under the behavior policy and then clones them [5, 26, 6, 17, 57].
Conditioned BC methods are based on the idea that every transition or trajectory is optimal when conditioned on the right variable [16, 11].
In this way, after conditioning, the data becomes optimal given the value of the conditioned variable.
The conditioned variable could be the cumulative return [8, 20, 28], or the goal information if provided [35, 16, 11, 60].
A recent study [11] shows that goal-conditioning can be super useful in D4RL AntMaze datasets.
However, this introduces the assumption that we have some prior information about the structure of the task, which is beyond the standard offline RL setup.
Our method can be viewed as a combination of RL-based and imitation-based methods.
We adopt an imitation-style manner to train the goal-policy and execute-policy.
During evaluation, however, following the guide-policy, the execute-policy could produce out-of-distribution actions by leveraging the generalization capacity of the function approximator.
Preliminaries Offline RL
We consider the standard fully observed Markov Decision Process (MDP) setting [47].
An MDP can be represented as M = ( S , A , T, r, ρ, γ ) where S is the state space, A is the action space, T ( ·| s, a ) is the transition probability distribution function, r ( s, a ) is the reward function, ρ is the initial state distribution and γ is the discount factor, we assume γ ∈ (0 , 1) in this work.
The goal of RL is to find a policy π ( a | s ) : S ×
A → [0 , 1] that maximizes the expected cumulative discounted reward (or called return) along a trajectory as
[Equation]
In this work, we focus on the offline setting.
The goal is to learn a policy from a fixed dataset D = (cid:8) τ i = (cid:0)(cid:0)
s i 0 , a i 0 , s ′
i 0 , r i 0 (cid:1) , (cid:0) s i 1 , a i 1 , s ′
i 1 , r i 1 (cid:1) , · · · , (cid:0) s iH , a iH ,
s ′ iH , r iH (cid:1)(cid:1)(cid:9) N i =1 consisting of trajectories that are collected by different policies, where H is the time horizon.
The dataset can be heterogenous and suboptimal, we denote the underlying behavior policy of D as µ , which represents the conditional distribution p ( a | s ) observed in the dataset.
RL via Supervised Learning
Conventional RL methods generally either compute the derivative of (1) with respect to the policy directly via policy gradient methods [44, 21], or estimate a value function or Q-function by means of TD learning [54], or both [45, 18].
In the RL via Supervised
Learning framework, we avoid the complex and potentially high-variance policy gradient estimators, as well as the complexity of temporal difference learning.
Instead, we perform conditioned behavior cloning on some extra information, such as goal state, cumulative return, or language description [35, 16].
When applying the framework to the offline RL setting, we can take the offline dataset D as input and find the outcome-conditioned policy π that optimizes
[Equation]
(2) where ω is an outcome conditioned on the remaining trajectory, i.e., ω ∼ g ( ·| τ t : H ) , and we use τ i : j = ( s i , a i , . . .
, s j ) to denote a fragment of the trajectory.
Policy-guided Offline Reinforcement Learning
We now continue to describe our proposed approach, POR .
To have a clear understanding of the out-of-distribution generalization of POR , we begin with a motivating example which demonstrates that by allowing state-compositionality on the dataset, the agent is able to generalize to a much better policy than doing action-compositionality.
Then we describe how we accomplish this idea in the more general case and introduce the algorithm details of POR .
Finally, we give a theoretical analysis of POR , we are interested to see how close the policy induced by POR and the optimal policy could be.
4.1 A Motivating Example
We first use a simple, didactic toy example to illustrate the importance of state-stitching for offline RL.
Let’s consider the navigation task shown in Figure 1, where the goal is to find the shortest path from the starting location (bottom-left corner) to the goal location (top-right corner) in the grid world.
The state and action space in this task are discrete, the state consists of the ( x, y ) coordinate and the agent could choose to move to any of its nearby 8 grids (i.e., the action space is 8).
The agent receives a reward of 0 for entering the goal state and − 1 for all other transitions.
The offline dataset (green-arrowed lines) consists of two suboptimal trajectories from the start location to the goal location, plus transitions of taking random actions at random states.
We use the grid color (more saturated to purple means higher value) to show the value function V ( s ) computed from the offline dataset.
We compare two methods that use the value function differently.
At a state s b ∼ D , the first method chooses the action a b in the dataset that leads to s ′ b with the highest V ( s ′ b ) ; the second method chooses the action a that leads to s ′ in the dataset with the highest V ( s ′ ) , note that a is not necessarily to be within the dataset.
The first method is doing action "stitching" in the dataset while the second method is doing state "stitching".
Although both methods are able to "stitch" suboptimal trajectories together from the dataset and successfully travel from the start location to the goal location, the actionstitching method takes more steps (11 steps) compared with the state-stitching method (7 or 8 steps), as shown in Figure 1.
The action-stitching method introduced above is actually an abstract methodology of imitation-based methods, which imitate the transitons in the dataset unequally with different weight choices.
However, these methods are too conservative to find the optimal trajectory, especially when there does not exist one in the offline dataset.
To generate better-than-dataset trajectories (e.g., red-arrowed lines), the agent needs to take logical out-of-distribution actions.
Note that in the toy example, we give the agent the freedom to choose any action that leads to the highest-value state in the dataset.
This allows the most out-of-distribution generalization and is realizable in some simple tasks (e.g., discrete state and action space) with high action coverage, for instance, in the toy example, given historical ( s, a, s ′ ) results at (5 , 4) , the agent can successfully recognize that taking upper-right will arrive at (6 , 5) .
However, in the more general case, especially the continuous state and action space setting, this may make the agent do erroneously-generalized actions as the same state can hardly be observed twice.
To perform logical out-of-distribution generalization, we need a Prophet , i.e., the guide-policy, to help the agent by telling it which state it should (high reward) and can (logical generalization) go to.
It turns out that the involvement of the guide-policy also brings some theoretical meaning, we will discuss that in Section 4.4.
4.2 Learning the Guide-Policy
Recall that the goal of the guide-policy is to guide the execute-policy about which state s it should go to.
To accomplish that, we train a state value function V ( s ) :
S → R .
The training of V uses only ( s, s ′ )
samples in the offline dataset, it doesn’t suffer from overestimation because there’re no out-of-distribution actions involved.
To approximate the optimal value function in the dataset, we adopt tricks from recent work [26, 36] by giving the ℓ 2 loss with a different weight using expectile regression, yielding the following asymmetric ℓ 2 loss min ϕ E ( s,r,s ′ ) ∼D [
[Equation]
2 ( r + γV ϕ
( s ′ ) − V ϕ ( s ))] , where L τ 2 ( u ) = | τ − 1 ( u < 0) | u 2 . It can be seen that when τ = 1 / 2 , this operator is reduced to Bellman expectation operator, while when τ → 1 , this operator approaches Bellman optimality operator.
Note that our learning objective bears similarity with IQL [26], however, here we aim to learn a state value function while IQL aims to learn a state-action value function so as to extract the policy.
To do this, IQL needs to learn both Q and V ( V is used to isolate the effect of state and approximate the expectile of Q only with respect to the action distribution), the training of V and Q are coupled and may affect each other.
Simply maximizing the guide-policy with respect to V ( s ) will result in a state where the executepolicy may make erroneous generalization.
To alleviate this issue, we add a behavior cloning term to the learning objective of the guide-policy.
Denote the guide-policy as g ω ( s ) : S → S , the learning objective of g is given by
[Equation]
where the weight α serves as the trade-off between guiding to space with high-reward and space that the execute-policy ought to have a correct generalization.
We also give an alternative learning objective that eliminates the need of α , by using the residual ( r + γV ϕ ′ ( s ′ )
− V ϕ ( s ) ) as the behavior cloning weight, which we find work well in practice:
[Equation]
4.3 Learning the Execute-Policy: Training and Evaluation
Algorithm
Require: D , τ , α (optional).
Training 2: Initialize V ϕ , V ϕ ′ , g ω , π θ 3: for t = 1 , 2 , · · · , N do 4:
Sample transitions ( s, r, s ′ ) ∼ D 5:
Update V ϕ by Eq.(3) 6:
Update g ω by Eq.(4) or Eq.(5) 7:
Update V ϕ ′ by ϕ ′ ← λϕ + (1 − λ ) ϕ ′ 8: end for 9: for t = 1 , 2 , · · · , M do 10:
Sample transitions ( s, a, s ′ ) ∼ D 11:
Update π θ by Eq.(6) 12: end for 13: //
Evaluation 14: Get initial state s , set d as False 15: while not d do 16:
Get action a form Eq.(7) 17:
Roll out a and get ( s ′ , r, d ) 18:
Set s = s ′ 19: end while
After learning the guide-policy, now we turn to the execute-policy.
Since the job of the execute-policy is to have a strong generalization ability, we adopt the RL via Supervised Learning framework by conditioning the execute-policy on s ′ that encountered in the dataset.
To be more specific, denote the executepolicy as π θ ( s, s ′ ) :
S × S → [0 , 1] .
During training, π performs supervised learning by maximizing the likelihood of the actions given the states and next states, yielding the following objective
[Equation]
in some scenarios with low data quality, we also add the exponential residual weight to (6) to mitigate the effect of those bad actions (same as (5)).
During evaluation, given a state s , the final action is determined by both the guide-policy and the executepolicy, by
[Equation]
Note that our learning objective (6) differs from previous RL via Supervised Learning methods in that we only use the next state s ′ as the conditioned variable, there’s no need to estimate the cumulative return [8], which may be highly suboptimal in certain cases [12, 4]. Our method also works in settings where we don’t know the goal information [11].
Also, during evaluation, previous methods show that the choice of conditioned variable is crucial important as little changes will cause significant performance difference [8].
In our method, owing to the existence of the guide-policy, the optimal conditioned variable can be automatically generated.
Our final algorithm, POR , consists of three supervised stages: learning V , learning g , and learning π .
We summarize the training and evaluation procedure of POR in Algorithm 1.
Note that the training of g and π are fully decoupled, this brings POR some nice properties to improve with supplementary data or transfer to new tasks, we will further investigate it in the experiments.
4.4 Analysis
In this section, we give a theoretical analysis of POR .
Concretely, we aim to 1) give the lower bound of the performance difference between POR and the optimal policy π ∗ , and 2) analyze how the guide-policy g will influence this bound.
We begin by introducing the notation and assumptions used in our analysis.
We denote P ( s, s ′ ) :
S × S →
A as the inverse transition operator .
We denote a g and a as the ground truth of π ( s, g ( s )) and π ( s, s ′ ) , respectively.
We also denote ϵ :=
sup ( s,a,s ′ ) ∈D ∥
π ( s, s ′ ) − a ∥ as the upper bound of the approximation error in the dataset during training.
We then introduce the following two assumptions.
Assumption 1.
(Non-lazy MDP)
The inverse transition operator P ( s, s ′ ) is L 1 -Lipschitz continuous, i.e., [Equation] P ( s 1 , s [Equation]
[Equation]
P (
[Equation]
, s [Equation] [Equation] [Equation]
1 , s
[Equation]
[Equation]
∥ .
Assumption 2.
(Lipschitz continuous function approximators)
The execute-policy we trained is
L 2 -Lipschitz continuous, i.e., ∥ π ( s 1 , s ′ 1 ) − π ( s 2 , s ′ 2 ) ∥
≤ L 2 ∥ ( s 1 , s ′ 1 ) − ( s 2 , s ′ 2 ) ∥ .
Assumption 1 holds when actions do have effects on state transitions (i.e. if a ̸ = 0 , then ∥ s ′ − s ∥ ≥ ϵ s > 0 ).
We refer the MDP that satisfies Assumption 1 as non-lazy MDP and others as lazy MDP .
It should be mentioned that most real-world tasks belong to the non-lazy MDP case, it is meaningless to study under the lazy MDP case because in lazy MDP , action changes will have little effect on state transitions, making the policy learning in vain.
Assumption 2 is a mild assumption that is frequently utilized in plenty of works [52, 37, 33].
Then in Theorem 1, we show that under these two assumptions, the gap between the optimal action and the action output by POR can be bounded.
Theorem 1.
(Single step gap to optimal action).
The single-step gap between optimal action and the action induced by our method can be bounded as
[Equation]
Theorem 1 states that single step optimal gap is related to three parts: l 1 , l 2 , l 3 .
l 1 is the generalization performance factor.
It can be found that a small value of ∥ g ( s ) − s ′ ∥ enables a small l 1 , which indicates the necessity to constrain the guide-policy to stay close to the dataset.
l 2 denotes the suboptimality constant.
Intuitively, a loosened constraint on the guide-policy may induce a small suboptimality constant but on the other hand, may suffer from the risk of a high value of l 1 .
The guide-policy g serves as a trade-off between l 1 and l 2 , we could balance the constraint strength and suboptimality by adjusting α in Eq.(4) to obtain the lowest bound.
l 3 is related to the approximation error on the training data, which can be reduced via improving the representation ability of function approximators but may suffer from a high L 2 due to the overfitting caused by over-parameterization. l 1 also depends on the Lipschitz constant L 1 and L 2 .
Note that a smooth inverse transition operator will result in a small value of L 1 and L 2 (because π ( s, s ′ ) is trained to approximate P ( s, s ′ ) ), this means our method may achieve better performance under a smoother inverse transition.
Based on Theorem 1, we can also give the performance bound of our method by replacing the ϵ in
Theorem 3.2 in [16] with the RHS of Eq.(8): J ( π ∗ ) − J ( π ) ≤ sup s,a,h ( l 1 + l 2 + l 3 ) T , where T is the horion of each episode, h is the time step.
Experiments
We present empirical evaluations of POR in this section.
We first evaluate POR against other baseline algorithms on D4RL [13] benchmark datasets.
We then explore deeper on the guide-policy about the benefits of the decoupled training process.
We finally establish ablation studies on the execute-policy.
5.1 Comparative Experiments on D4RL Benchmark Datasets
We first evaluate our approach on D4RL MuJoCo and AntMaze datasets [13].
Notice that most of MuJoCo datasets (except random and medium datasets) consist of a large fraction of near-optimal trajectories, while AntMaze datasets contain few or no near-optimal trajectories.
Hence those low-quality datasets can serve as a good testbed for quantifying the compositionality ability.
We compare POR with both RL-based and imitation-based baselines.
For RL-based baselines, we select the best one from 5 state-of-the-art algorithms, including BCQ [15], BEAR [27], BRAC [55],
CQL [29] and TD3+BC [14].
For imitation-based baselines, we further split them into Weighted
BC methods and Conditioned BC methods.
Weighted BC methods include 10%BC [8], One-step
RL [5] and IQL [25].
10%BC is a filtered version of BC that runs behavior cloning on only the top 10% high-return trajectories in the dataset.
One-step RL and IQL can be viewed as using different weights to do behavior cloning (see Appendix A for discussion).
Conditioned BC methods include
DT [8] and RvS [11].
DT is built with Transformer and conditioned on the cumulative return.
RvS well studied the effect of different conditional information on different tasks.
The score of RvS is [Table] chosen by the higher score between RvS-R and Rvs-G, which use the reward-to-go and oracle goal information as the conditioned variable, respectively.
Note that we also categorize our method POR into conditioned BC methods because the execute-policy is learned in a similar manner to other conditioned BC methods.
Note that the results of Filter BC are from our own implementation, and the results of IQL on MuJoCo random tasks are re-runed using the author-provided implementation.
Other results are taken directly from their corresponding papers.
To enable fairness and consistency in the evaluation process, we train our algorithm for 1 million time steps and evaluate everytime steps that consist of 10 episodes.
Full experimental details are included in Appendix D.1 and the learning curve can be found in Appendix E.
The results are shown in Table 1, in MuJoCo locomotion tasks, POR performs competitively to the best performance of prior methods in high-quality datasets, while achieving much better performance than other methods in low-quality datasets (e.g., random and medium datasets).
On the more challenging
AntMaze tasks, POR outperforms all other baselines by a large margin.
It even surpassed RvS that uses the oracle goal information to learn the policy.
We conjecture that the success in all those datasets should be credited to the out-of-distribution generalization ability of POR .
5.2 Validation Experiments on Out-of-Distribution Generalization
In this section, we try to validate our hypothesis about the out-of-distribution generalization of POR .
We are especially interested to see how the guide-policy is learned in those high-dimensional tasks in D4RL.
To do so, we compare the distribution of g ( s ) generated by the guide-policy with the distribution of s ′ from the offline dataset.
We choose antmaze-large-diverse and hopper-medium-replay , two datasets for evaluation.
For each dataset, we randomly collect 100 , 000 ( s, s ′ ) transition tuples.
To visualize the results clearly, we plot the distribution of g ( s ) and s ′ with t-Distributed Stochastic Neighbor Embedding (t-SNE) [19], we also visualize the true values of all states using the optimal value function V ∗ ( s ) . [Figure]
The results are shown in Figure 2, it can be found that in both tasks, the distribution of g ( s ) bears some similarity to the distribution of s ′ .
However, when looking at the non-overlapped area of g ( s ) and s ′ , we found that g ( s ) has a much higher value than s ′ , especially in the antmaze-large-diverse dataset.
Those out-of-distribution and high-value areas are exactly where POR generalized to, note that the high-value benefit of the single-step generalization could also be accumulated and propagated through time steps, resulting in a much better policy.
5.3 Investigations on the Guide-Policy
In this section we investigate what is the potential benefit of decoupling the training process of the guide-policy and execute-policy.
We found that as the execute-policy is not task-specific and only cares about the effect of actions on different states, one can reuse its powerful generalization ability when encountered with additional suboptimal data or adapting to a new task.
We only need to re-train the guide-policy as it is task-specific.
Improve guide-policy by additional suboptimal data
We first study the setting where we already have a policy learned from a small yet high-quality dataset D e , and now we get a supplementary large dataset D o .
D o may be sampled from one or multiple behavior policies, it provides higher data coverage but could be suboptimal.
We are interested to see whether we can use D o to obtain a better policy.
This setting is realistic as low-quality datasets appear more often in real-world tasks [61, 58].
Naïvely, we can combine D e and D o as a new offline dataset D , then run any off-the-shelf offline
RL algorithm.
We label this training scheme as "more".
We also label the original training scheme that only uses D e as "main".
We test these two training schemes on two state-of-the-art offline RL algorithms, TD3+BC and CQL.
The results shown in Figure 3 indicate that both "main" and "more" [Figure] are insufficient to learn a good policy, The reason why "main" fails is that the resulting policy may not be able to generalize due to the limited size.
While "more" can generalize better than "main" due to access to a much larger dataset, the policy will be negatively impacted by the low-quality data in D o as both TD3+BC and CQL contain the behavior regularization term.
Due to the availability of decoupled training, POR is able to use a different training scheme, we can keep the execute-policy learned from D e unchanged but use the combination of D e and D o to re-train the guide-policy.
Note that in this scheme, we allow the use of an action-free dataset D o , which often appears in real-world scenarios (i.e., video data and third-person demonstrations) [50].
The value function will be more accurate and better generalized when combining D e with D o to have a larger data coverage (see exploration in online RL as an example).
We label this unique training procedure as "mix".
It can be shown in Figure 3 that, by doing so, POR could have a super large improvement over "main" and "more".
The version "mix" brings POR up to 99 , close to the upper limit of the performance in the hopper environment, with only limited data.
In the walker2d environment, POR with "mix" also outperforms all other baselines.
Change guide-policy in new tasks
Then, we study whether we can use the guide-policy to do a lightweight adaptation to new tasks, without changing the execute-policy.
We conduct experiments on the continuous variant of the classic four-room environment [20].
The training data consists of trajectories collected by a goal-reaching controller with the start and end locations sampled randomly at non-wall locations, mixed with data collected by a random policy.
In this environment, we introduce three different tasks: Task A ( Four-room ) needs the agent to travel from the start location to the goal location, besides that,
Task B ( Four-room-river ) and task C
( Four-room-key ) require the agent to bypass the river or get a key before arriving at the goal location.
We first train the guide-policy and the execute-policy π to solve task A.
In task B and task C, we re-train the guide-policy using the corresponding reward function in that task and reuse the executepolicy π from Task A.
In Figure 4, we can see that by doing so, POR is able to successfully solve task
B and task C.
In detail, in task B, POR is aware of the river and correctly bypasses it.
In task C, POR only chooses to pass the bottom door, this is owing to the awareness of the importance of obtaining the key before reaching the goal.
This result shows the strong adaptation ablity of the guide-policy for different tasks, as well as the generalization ability of the execute-policy across different tasks.
Conclusions and Limitations
In this work, we propose a novel offline RL approach, POR , which leverages the training stability of imitation-style methods while still encouraging logical out-of-distribution generalization.
POR allows state-compositionality rather than action-compositionality from the dataset.
Through theoretical analysis and extensive experiments, we show that POR outperforms prior methods on a variety of datasets, especially those with low quality.
We also empirically demonstrate the additional benefits of
POR in terms of improving with supplementary suboptimal data and easily adapting to new tasks.
One limitation of POR is that the prediction error of the guide-policy may be large when the state space is high-dimensional.
However, this can be alleviated with the help of representation learning [31, 65]
Acknowledgments
This work is supported by fundings from AsiaInfo.
We thank Junwu Xiong,
James Zhang for their feedback on earlier versions of the manuscript
A More Discussion Why One-step and IQL are imitation-based methods?
The core difference between RL-based and imitation-based methods is that RL-based methods learn a value function of policy π while imitation-based methods don’t.
Learning the value function of π requires off-policy evaluation of π (i.e., learning Q π or V π ), which is prone to distribution shift.
The policy evaluation and policy improvement will also affect each other as they are coupled.
Imitation-based methods don’t learn Q π or V π , but some of them do learn a value function.
We call these methods imitation-based because they learn the value function using only dataset samples, the value function actually tells how advantageous it could be under the behavior policy, much like imitation learning.
Also, the policy learning objective of One-step (with exponentially-weighted improvement operator) and IQL can be written as L ( π ) =
E ( s,a ) ∼D [exp ( β ( Q ( s, a ) − V ( s ))) log π ( a | s )] , which uses dataset actions to doing behavior cloning with different weights (One-step and IQL learn different value functions).
Cloning dataset actions can only do action-stitching, which loses the ability to surpass the dataset by out-of-distribution (action) generalization.
For example, in our toy example in Section 4.1, actionstitching methods at most learn the shortest path that contained in the datset.
This is suboptimal especially when there does not exist one complete path starting from the start location to the goal location in the offline dataset.
How can we do beyond action-stitching?
One way is to use RL-based methods, by querying an accurate Q π , we can get a different-yet-optimal action a π , but Q π is hard to estimate.
Another way is like what POR did, we learn an out-of-distribution state indicator, i.e., the guide-policy g , to guide the policy to the optimal next state.
If the execute-policy can generalize well, it will also output a different yet optimal action a = arg max a π ( a | s, g ( s )) .
B More Related Work
Our work decouples the state-to-action policy into two modules, i.e., the guide-policy and the executepolicy.
The execute-policy is actually an inverse dynamics model, which has been widely used in various ways in sequential decision-making.
In exploration, inverse dynamics can be used to learn representations of the controllable aspects of the state [40].
In imitation learning, [50] and [34] train an inverse dynamics model to label the state-only demonstrations with inferred actions.
[10] use inverse dynamics models to translate actions taken in a simulated environment to the real world.
Recently, there has been an emergence of work [46, 16] highlighting the connection of imitation learning and reinforcement learning.
Specifically, rather than learn to map states and actions to reward, as is typical in reinforcement learning, [46] trains a model to predict actions given a state and an outcome, which could be the amount of reward the agent is to collect within a certain amount of time.
[16] uses a similar idea, predicting actions conditioned on an initial state, a goal state, and the amount of time left to achieve the goal.
These methods are perhaps the closest work to our algorithm, however, we study the offline setting and motivate the usage of an inverse dynamics model from a different perspective (i.e., state-stitching).
C Proof
In this section, we provide the proof of Theorem 1.
Proof.
We can rewrite the LHS of Eq.(8) as
[Equation]
a
[Equation]
a
[Equation]
− a
[Equation]
∥ ≤
[Equation]
[Equation]
[Equation]
a
[Equation]
a
[Equation]
a
[Equation]
(Triangle) ≤ L
[Equation]
g ( s ) − s
[Equation]
+
[Equation]
+ L
[Equation]
g ( s ) − s ′ ∥
[Equation]
∥ a g − a ∗ ∥ (Assumption1&2) ≤ ( L 1 + L 2 )
∥ g ( s ) − s ′ ∥
[Equation]
D Experimental Details
In this section, we provide the experimental details of our paper.
We use the following hardware and software for our training:
• GPUs: NVIDIA GeForce RTX 3080Ti • Python 3. • Tensorflow 1.15.0 [1] • Gym {+0.23.
 [7] •+} MuJoCo {+2.1.
 [49] •+} mujoco-py [-2.1.2.
-]2.1.2.
D.1 D4RL Experiments Data collection
The datasets in D4RL have been generated as follows:
random : roll out a randomly initialized policy for 1M steps.
expert : 1M samples from a policy trained to completion with SAC [18].
medium : 1M samples from a policy trained to approximately 1/3 the performance of the expert. medium-replay : replay buffer of a policy trained up to the performance of the medium agent. medium-expert : 50-50 split of medium and expert data.
For all datasets we use the v2 version.
Implementation details
Our implementation of 10%BC is as follows, we first filter the top 10 % trajectories in terms of the trajectory return, and then run behaviour cloning on those filtered data.
The hyperparameters of POR are present in Table 2.
We use target networks for the V -function and use clipped double V -learning (take the minimum of two V -functions) for all updates.
We normalize state to [ − 1 , 1] [43] to reduce the prediction error of the guide-policy, it can be deemed as an naïve method of representation learning.
[Table]
POR Hyperparameters.
[Table]
[Table]
D.2 Additional Suboptimal Data Experiments Data collection and settings
In this experiment, we use different part of medium-replay datasets as D e and D o .
More specific, we use 30% to 70% transitions to constitute D e and use 20% to 80% transitions to constitute D e ∪D o .
Implementation details
We use the same hyperparamters of POR , shown in Table 2.
Our implementations of TD3+BC † [15], CQL ‡ [29] is from the author-provided implenmentation from Github, and we keep all parameters the same to the author-provided implementation.
[Table] [Table]
D.3 Four-room Experiments Environment settings
We use the continuous variant of the classic four-room environment from [20].
This continuous variant of four-rooms is basically the same as the traditional classic four-room environment in the environment design.
There are 19 × 19 grids that consist of four rooms with only one block channel with neighboring rooms.
The goal of this environment is to make the agent travel from one location to another different location.
Those tasks are challenging as their reward is extremely spare, they need the agent to have the ability to explore efficiently through the whole state space.
The action and observation space are shown in Table 5.
The first and second dimension of the action space represents the distance to travel on the x -axis and y -axis, respectively.
The first and second dimension of the observation space represents the coordinate on the x -axis and y -axis, respectively.
[Table]
In task A ( Four-room ), only reaching the goal will give the agent a reward of 1 , otherwise the agent will get 0 reward.
When the agent reaches the target or the agent takes over 500 steps, the environment will terminate.
In task B ( Four-room-river ), the agent will receive − 1 reward if it falls into the river, and the environment will be terminated.
Also, the agent will get 1 reward only when it reaches the goal.
In task C ( Four-room-key ), the agent gets 1 reward only when it reaches the goal and gets the key.
Falling into the river will also give the agent − 1 reward.
Data collection and implementation details
The training data consists of trajectories collected by a goal-reaching controller with the start and end locations sampled randomly at non-wall locations.
To make the data more diverse, we also add trajectories from a random policy.
We collect 100, transitions for each task.
Figure 6 shows the training hyperparameters of POR in the three four-room environments.
[Table]
E Learning Curves
In this section
[Table]
we provide the learning curves of our experiments in the main paper
[Table]
[Figure] [Figure]
N o r m a li
z ed S c
o r e
[Figure]
N o r m a li
z ed S c
o r e
[Figure]
Iterations
[Figure] [Figure]
[Figure] [Figure]
F Ablation Study on the Execute-Policy
We also present an ablation study on the execute-policy, with the aim to answer the following two questions: 1) How does network capacity affect the performance of the execute-policy? 2)
How does the performance change if we use partial "good" data, instead of full data, to train the execute policy?
To answer the first question [Figure] we compare two choices of network size: a big network with (512 , 512) hidden units and a small network with (128 , 128) hidden units.
To answer the second question [Figure] we compare with the choice that only selects a subset of the dataset to train the execute policy.
Concretely, we choose the top X % trajectories in the dataset, ordered by episode returns.
We sweep X over [10 , 25 , 40] and choose the best score.
We give the mean scores of AntMaze ( A ) and MuJoCo ( M ) datasets in Table 7.
It can be seen that adopting a big network consistently gives a better performance, which is also found in [11].
Using partial data will result in a less-performed policy, especially on AntMaze datasets.
Note that in MuJoCo datasets, the performance didn’t drop too much.
This is because MuJoCo datasets don’t require the compositionality ability, using partial trajectories could already achieve high scores.
[Table]