OfflineRL-Kit: 一个优雅的 PyTorch 离线强化学习库

引言

在机器学习和人工智能领域，强化学习一直是一个备受关注的研究方向。然而，传统的在线强化学习方法在某些场景下可能存在局限性，例如在真实世界中收集大量数据可能成本高昂或危险。为了解决这一问题，离线强化学习应运而生。OfflineRL-Kit 正是为此而设计的一个强大工具，它为研究人员提供了一个优雅、高效的离线强化学习开发环境。

OfflineRL-Kit Logo

OfflineRL-Kit 简介

OfflineRL-Kit 是一个基于纯 PyTorch 的离线强化学习库，由 Yihao Sun 开发。这个库的设计理念是为研究人员提供一个友好、便捷的工具，帮助他们更快速地进行离线强化学习算法的开发和实验。

主要特性

优雅的框架：OfflineRL-Kit 的代码结构清晰，易于使用，让研究人员能够快速上手。
最先进的算法：库中包含了多种state-of-the-art的离线强化学习算法，涵盖了无模型和基于模型的方法。
高可扩展性：研究人员可以基于库中的组件，用很少的代码就能构建新的算法。
并行调优支持：方便研究人员进行大规模的参数调优实验。
强大的日志系统：清晰而功能强大的日志系统，便于管理和分析实验结果。

支持的算法

OfflineRL-Kit 支持多种先进的离线强化学习算法，大致可以分为无模型方法和基于模型的方法两类：

无模型方法

Conservative Q-Learning (CQL)：一种保守的 Q 学习方法，通过最小化某些动作的 Q 值来避免对未见过的状态-动作对的过度估计。
TD3+BC：结合了 TD3 算法和行为克隆（Behavior Cloning）的方法，在离线设置中表现良好。
Implicit Q-Learning (IQL)：一种隐式 Q 学习方法，通过学习值函数的分位数来改善策略学习。
Ensemble-Diversified Actor Critic (EDAC)：使用集成方法来增强策略的多样性，提高离线学习的鲁棒性。
Mildly Conservative Q-Learning (MCQ)：相较于 CQL，这是一种较为温和的保守 Q 学习方法。

基于模型的方法

Model-based Offline Policy Optimization (MOPO)：通过学习环境模型并在模拟环境中进行策略优化，同时考虑模型不确定性。
Conservative Offline Model-Based Policy Optimization (COMBO)：在 MOPO 的基础上增加了保守性约束，进一步提高了性能。
Robust Adversarial Model-Based Offline Reinforcement Learning (RAMBO)：引入对抗性训练来增强模型的鲁棒性。
Model-Bellman Inconsistency Penalized Offline Reinforcement Learning (MOBILE)：通过惩罚模型与贝尔曼方程的不一致性来改善离线学习。

安装指南

要使用 OfflineRL-Kit，您需要按照以下步骤进行安装：

首先，安装 MuJoCo 引擎。您可以从 MuJoCo 官网下载。安装完成后，根据您安装的 MuJoCo 版本安装相应的 mujoco-py。
接下来，安装 D4RL：

git clone https://github.com/Farama-Foundation/d4rl.git
cd d4rl
pip install -e .

最后，安装 OfflineRL-Kit：

git clone https://github.com/yihaosun1124/OfflineRL-Kit.git
cd OfflineRL-Kit
python setup.py install

快速开始

为了帮助您快速上手 OfflineRL-Kit，这里提供了一个使用 Conservative Q-Learning (CQL) 算法的简单示例。

训练过程

首先，创建环境并获取离线数据集：

env = gym.make(args.task)
dataset = qlearning_dataset(env)
buffer = ReplayBuffer(
buffer_size=len(dataset["observations"]),
obs_shape=args.obs_shape,
obs_dtype=np.float32,
action_dim=args.action_dim,
action_dtype=np.float32,
device=args.device
)
buffer.load_dataset(dataset)

定义模型和优化器：

actor_backbone = MLP(input_dim=np.prod(args.obs_shape), hidden_dims=args.hidden_dims)
critic1_backbone = MLP(input_dim=np.prod(args.obs_shape) + args.action_dim, hidden_dims=args.hidden_dims)
critic2_backbone = MLP(input_dim=np.prod(args.obs_shape) + args.action_dim, hidden_dims=args.hidden_dims)
dist = TanhDiagGaussian(
    latent_dim=getattr(actor_backbone, "output_dim"),
    output_dim=args.action_dim,
    unbounded=True,
    conditioned_sigma=True
)
actor = ActorProb(actor_backbone, dist, args.device)
critic1 = Critic(critic1_backbone, args.device)
critic2 = Critic(critic2_backbone, args.device)
actor_optim = torch.optim.Adam(actor.parameters(), lr=args.actor_lr)
critic1_optim = torch.optim.Adam(critic1.parameters(), lr=args.critic_lr)
critic2_optim = torch.optim.Adam(critic2.parameters(), lr=args.critic_lr)

设置策略：

policy = CQLPolicy(
    actor,
    critic1,
    critic2,
    actor_optim,
    critic1_optim,
    critic2_optim,
    action_space=env.action_space,
tau=args.tau,
gamma=args.gamma,
alpha=alpha,
cql_weight=args.cql_weight,
temperature=args.temperature,
max_q_backup=args.max_q_backup,
deterministic_backup=args.deterministic_backup,
with_lagrange=args.with_lagrange,
lagrange_threshold=args.lagrange_threshold,
cql_alpha_lr=args.cql_alpha_lr,
num_repeart_actions=args.num_repeat_actions
)

定义日志记录器：

log_dirs = make_log_dirs(args.task, args.algo_name, args.seed, vars(args))
output_config = {
    "consoleout_backup": "stdout",
    "policy_training_progress": "csv",
    "tb": "tensorboard"
}
logger = Logger(log_dirs, output_config)
logger.log_hyperparameters(vars(args))

将所有组件加载到训练器中并开始训练：

policy_trainer = MFPolicyTrainer(
    policy=policy,
    eval_env=env,
    buffer=buffer,
    logger=logger,
    epoch=args.epoch,
    step_per_epoch=args.step_per_epoch,
    batch_size=args.batch_size,
    eval_episodes=args.eval_episodes
)

policy_trainer.train()

参数调优

OfflineRL-Kit 支持使用 Ray 进行并行参数调优。以下是一个简单的示例：

ray.init()
# 加载默认参数
args = get_args()

config = {}
real_ratios = [0.05, 0.5]
seeds = list(range(2))
config["real_ratio"] = tune.grid_search(real_ratios)
config["seed"] = tune.grid_search(seeds)

analysis = tune.run(
    run_exp,
    name="tune_mopo",
    config=config,
    resources_per_trial={
        "gpu": 0.5
    }
)

日志系统

OfflineRL-Kit 的日志系统支持多种记录文件类型，包括：

.txt（控制台输出备份）
.csv（记录训练过程中的损失、性能或其他指标）
.tfevents（用于 TensorBoard 可视化训练曲线）
.json（超参数备份）

日志系统的结构清晰，便于管理实验：

└─log(root dir)
    └─task
        └─algo_0
        |   └─seed_0&timestamp_xxx
        |   |   ├─checkpoint
        |   |   ├─model
        |   |   ├─record
        |   |   │  ├─tb
        |   |   │  ├─consoleout_backup.txt
        |   |   │  ├─policy_training_progress.csv
        |   |   │  ├─hyper_param.json
        |   |   ├─result
        |   └─seed_1&timestamp_xxx
        └─algo_1

使用日志记录器的示例：

from offlinerlkit.utils.logger import Logger, make_log_dirs

log_dirs = make_log_dirs(args.task, args.algo_name, args.seed, vars(args))
output_config = {
    "consoleout_backup": "stdout",
    "policy_training_progress": "csv",
    "dynamics_training_progress": "csv",
    "tb": "tensorboard"
}
logger = Logger(log_dirs, output_config)
logger.log_hyperparameters(vars(args))

# 记录一些指标
logger.logkv("eval/normalized_episode_reward", norm_ep_rew_mean)
logger.logkv("eval/normalized_episode_reward_std", norm_ep_rew_std)
logger.logkv("eval/episode_length", ep_length_mean)
logger.logkv("eval/episode_length_std", ep_length_std)
# 设置时间步
logger.set_timestep(num_timesteps)
# 将结果写入记录文件
logger.dumpkvs()

结果可视化

OfflineRL-Kit 提供了简单的脚本来绘制实验结果：

python run_example/plotter.py --algos "mopo" "cql" --task "hopper-medium-replay-v2"

这将生成指定算法和任务的性能对比图。

总结

OfflineRL-Kit 是一个功能强大、易于使用的离线强化学习库，为研究人员提供了一个理想的实验平台。通过其优雅的框架设计、丰富的算法支持、高度的可扩展性以及强大的日志系统，研究人员可以更加高效地进行离线强化学习的研究和开发。

无论您是刚开始接触离线强化学习的新手，还是在该领域深耕多年的专家，OfflineRL-Kit 都能为您提供宝贵的支持。我们期待看到更多研究人员利用这个工具，在离线强化学习领域取得突破性的进展。

如果您在您的工作中使用了 OfflineRL-Kit，请引用以下 bibtex：

@misc{offinerlkit,
  author = {Yihao Sun},
  title = {OfflineRL-Kit: An Elegant PyTorch Offline Reinforcement Learning Library},
  year = {2023},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/yihaosun1124/OfflineRL-Kit}},
}

让我们一起探索离线强化学习的无限可能吧！🚀🤖🧠