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

OfflineRL-Kit是一个基于纯PyTorch的离线强化学习库。该库具有一些对研究人员友好且便利的特性，包括：

• 优雅的框架，代码结构清晰易用
• 最先进的离线强化学习算法，包括无模型和基于模型的方法
• 高度可扩展性，您可以基于我们库中的组件用几行代码构建新算法
• 支持并行调优，对研究人员非常方便
• 清晰强大的日志系统，易于管理实验

支持的算法

• 无模型
  • 保守Q学习 (CQL)
  • TD3+BC
  • 隐式Q学习 (IQL)
  • 集成多样化演员评论家 (EDAC)
  • 温和保守Q学习 (MCQ)
• 基于模型
  • 基于模型的离线策略优化 (MOPO)
  • 保守离线基于模型的策略优化 (COMBO)
  • 鲁棒对抗性基于模型的离线强化学习 (RAMBO)
  • 模型贝尔曼不一致性惩罚离线强化学习 (MOBILE)

基准结果（4个种子）（进行中）

详细日志可以在https://drive.google.com/drive/folders/11QHHDlLmUEc097tPgYvb4gZ2IaqSpkHp?usp=share_link中查看。

安装

首先，安装MuJuCo引擎，可以从这里下载，并安装mujoco-py（其版本取决于您安装的MuJoCo引擎版本）。

其次，安装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

快速开始

训练

这是一个CQL的示例。您也可以在run_example/run_cql.py运行完整脚本。

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

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()

调优

您可以借助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
    }
)

您可以在 tune_example/tune_mopo.py 查看完整脚本。

日志

我们的日志记录器支持多种记录文件类型，包括 .txt（用于备份标准输出）、.csv（记录训练过程中的损失或性能或其他指标）、.tfevents（用于可视化训练曲线的 tensorboard）、.json（用于备份超参数）。 我们的日志记录器还具有清晰的日志结构：

└─log(根目录)
    └─任务
        └─算法_0
        |   └─种子_0&时间戳_xxx
        |   |   ├─检查点
        |   |   ├─模型
        |   |   ├─记录
        |   |   │  ├─tb
        |   |   │  ├─控制台输出备份.txt
        |   |   │  ├─策略训练进度.csv
        |   |   │  ├─超参数.json
        |   |   ├─结果
        |   └─种子_1&时间戳_xxx
        └─算法_1

这是一个日志记录器的示例，您可以在 offlinerlkit/policy_trainer/mb_policy_trainer.py 查看完整脚本。

首先，导入一些相关的包：

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()

绘图

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

引用 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}},
}