This repo is the official implementation of Harnessing Mixed Offline Reinforcement Learning Datasets via Trajectory Weighting.

Abstract: Most offline reinforcement learning (RL) algorithms return a target policy maximizing a trade-off between (1) the expected performance gain over the behavior policy that collected the dataset, and (2) the risk stemming from the out-of-distribution-ness of the induced state-action occupancy. It follows that the performance of the target policy is strongly related to the performance of the behavior policy and, thus, the trajectory return distribution of the dataset. We show that in mixed datasets consisting of mostly low-return trajectories and minor high-return trajectories, state-of-the-art offline RL algorithms are overly restrained by low-return trajectories and fail to exploit high-performing trajectories to the fullest. To overcome this issue, we show that, in deterministic MDPs with stochastic initial states, the dataset sampling can be re-weighted to induce an artificial dataset whose behavior policy has a higher return. This re-weighted sampling strategy may be combined with any offline RL algorithm. We further analyze that the opportunity for performance improvement over the behavior policy correlates with the positive-sided variance of the returns of the trajectories in the dataset. We empirically show that while CQL, IQL, and TD3+BC achieve only a part of this potential policy improvement, these same algorithms combined with our reweighted sampling strategy fully exploit the dataset. Furthermore, we empirically demonstrate that, despite its theoretical limitation, the approach may still be efficient in stochastic environments.

• 2024/02/14: Added an example application to the newer version of d3rlpy codebase
• 2023/07/27: Added implementation on official codebase of CQL, IQL, and TD3+BC and their results
• 2023/06/15: Fix max-min normalization in AW and RW

We implemented both RW and AW in the newerd3rlpy codebase. See the example scripts here: d3rlpy_examples/train_cql_cartpole_weighted_sample.py. With the newer APIs, RW and AW can be applied easily with a dataset wrapper.

See d3rlpy_impls in this repo.

• Download and install suboptimal_offline_datasets
• Please refer to JaxCQL, implicit_q_learning, and td3bc directoriees.

Run python experiments/d4rl/run_cql.py --mode {screen,sbatch,local} --n_gpus {list of gpu ids} --n_jobs {number of parallel jobs}

Run python experiments/d4rl/run_iql.py --mode {screen,sbatch,local} --n_gpus {list of gpu ids} --n_jobs {number of parallel jobs}

Run python experiments/d4rl/run_td3bc.py --mode {screen,sbatch,local} --n_gpus {list of gpu ids} --n_jobs {number of parallel jobs}