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

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

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

支持的算法

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

	CQL	TD3+BC	EDAC	IQL	MOPO	RAMBO	COMBO	MOBILE
halfcheetah-medium-v2	49.4±0.2	48.2±0.5	66.4±1.1	47.4±0.5	72.4±4.2	78.7±1.1	71.9±8.5	75.8±0.8
hopper-medium-v2	59.1±4.1	60.8±3.4	101.8±0.2	65.7±8.1	62.8±38.1	82.1±38.0	84.7±9.3	103.6±1.0
walker2d-medium-v2	83.6±0.5	84.4±2.1	93.3±0.8	81.1±2.6	84.1±3.2	86.1±1.0	83.9±2.0	88.3±2.5
halfcheetah-medium-replay-v2	47.0±0.3	45.0±0.5	62.3±1.4	44.2±0.6	72.1±3.8	68.5±3.6	66.5±6.5	71.9±3.2
hopper-medium-replay-v2	98.6±1.5	67.3±13.2	101.5±0.1	94.8±6.7	92.7±20.7	93.4±11.4	90.1±25.2	105.1±1.3
walker2d-medium-replay-v2	71.3±17.9	83.4±7.0	86.2±1.2	77.3±11.0	85.9±5.3	73.7±6.5	89.4±6.4	90.5±1.7
halfcheetah-medium-expert-v2	93.0±2.2	90.7±2.7	101.8±8.4	88.0±2.8	83.6±12.5	98.8±4.3	98.2±0.2	100.9±1.5
hopper-medium-expert-v2	111.4±0.5	91.4±11.3	110.5±0.3	106.2±5.6	74.6±44.2	85.0±30.7	108.8±2.6	112.5±0.2
walker2d-medium-expert-v2	109.8±0.5	110.2±0.3	113.6±0.3	108.3±2.6	108.2±4.3	78.4±45.4	110.0±0.2	114.5±2.2

详细日志可以在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}},
}