AQuaDem 源码实战：基于演示动作量化的连续控制算法解析与运行指南

人工智能深度学习NLP计算机视觉强化学习【免费下载链接】google-researchGoogle Research项目地址https://gitcode.com/gh_mirrors/go/google-research点击查看免费下载AQuaDemActionQuantizationandDemonstrations全称 Continuous Control with Action Quantization from Demonstrations是 Google Research 团队提出的从演示中学习连续控制任务的算法它先把连续动作空间量化为一组由多模态行为克隆Multi-head BC学到的离散候选动作再在这些候选动作上运行标准离散强化学习算法进行决策。本文以 aquadem 目录源码为主线完整介绍该算法的依赖安装、运行命令、核心模块结构与配置参数并深入 learning.py、networks.py、builder.py 等源码解释其多分类 BC 预训练 离散 RL 微调的两阶段机制。读完本文你将能够独立安装并运行 AQuaDem理解每个命令行参数与配置文件项的含义并掌握将其迁移到自定义连续控制任务的要点。一、AQuaDem 算法思想概览AQuaDem 解决的核心问题是如何在连续动作空间上利用专家演示提升强化学习RL的训练效率与最终性能。直接在高维连续动作空间做 RL 通常样本效率低而纯粹的行为克隆BC又受限于演示质量、无法超越专家。AQuaDem 的做法是折中动作量化训练一个多模态动作编码器multi-modal encoder即多分类 BC给定观测时同时输出num_actions个候选连续动作覆盖专家演示中出现的多种行为模式离散化决策把原始连续动作空间替换为这N个候选动作环境的原始连续动作 spec 被改写为离散 spec见 discretize_spec离散 RL在候选动作上运行标准的离散动作 RL 算法本实现使用基于 Munchausen Q-Learning 的 DQN让 agent 学会选哪一个候选动作从而规避连续动作空间探索的困难。该目录为论文Continuous Control with Action Quantization from DemonstrationsRobert Dadashi、Leonard Hussenot 等arXiv 2110.10149的官方源码实现基于 DeepMind Acme 框架构建采用 JAX 实现网络与训练逻辑。二、环境准备与依赖安装AQuaDem 的依赖通过 requirements.txt 管理核心依赖包括Acme 生态dm-acme0.3.0、dm-reverb0.6.1回放缓冲区、dm-env、dm-haiku、dm-sonnet深度学习jax0.2.28、flax0.4.0、optax0.1.1、tensorflow2.7.0环境与数据gym0.21.0、dm-control、tensorflow_datasets4.4.0、rlds0.1.3以及通过 git 安装的rlaxDeepMind与d4rlBerkeley提供 Adroit 演示数据集。安装命令pip install -r requirements.txt两个需要特别注意的前提条件Python 版本官方声明 AQuaDem 兼容Python 3.9且依赖版本如jax0.2.28、tensorflow2.7.0均为较早期版本建议使用 Python 3.9 的虚拟环境安装避免与新版依赖产生冲突MuJoCo 物理引擎AQuaDem 需要在 MuJoCo2.1.1下运行Adroit 灵巧手任务依赖 MuJoCo请按 MuJoCo 官方提供的 2.1.1 版本安装指引完成安装后再运行上面的pip install。三、快速启动运行 AQuaDem安装完成后在仓库根目录执行以下命令即可启动训练python -m aquadem.run_aquadqn --workdir/tmp/aquadem --env_namedoor-human-v1其中workdir为日志输出目录训练与评估指标以 CSV 形式写入env_name为要运行的环境。默认情况下会使用 100 万步环境交互训练door-human-v1Adroit 门开关任务的人类演示版本。3.1 命令行参数入口脚本 run_aquadqn.py 通过 absl.flags 定义全部命令行参数参数类型默认值说明--workdirstr/tmp/aquadqn日志输出目录训练/评估指标以 CSV 保存--env_namestrdoor-human-v1运行的环境名称D4RL Adroit 任务如door-human-v1、hammer-human-v1等--num_demonstrationsintNone使用的专家演示条数None表示使用完整数据集--num_stepsint1000000训练总环境步数--eval_everyint10000评估频率每多少步评估一次--seedint0RL agent 的随机种子注意源码要求num_steps % eval_every 0run_aquadqn.py否则程序会通过 assert 中断。3.2 训练-评估循环run_aquadqn.py 中采用先评估后训练的交替循环for _ in range(FLAGS.num_steps // FLAGS.eval_every): eval_loop.run(num_episodes10) # 每次评估跑 10 个 episode train_loop.run(num_stepsFLAGS.eval_every) # 再训练 eval_every 步 eval_loop.run(num_episodes10)评估环境在训练环境基础上额外包裹了SuccessRewardWrapperwrappers.py即整个 episode 首次累积回报达到阈值即返回奖励 1因此评估回报就是是否成功的 0/1 指示便于直观衡量任务成功率。四、源码架构总览aquadem目录共有 9 个文件职责划分清晰文件职责run_aquadqn.py程序入口组装环境、builder、网络、训练/评估循环config.pyAquademConfig数据类集中管理全部算法超参数builder.pyAquademBuilder实现 Acme 的ActorLearnerBuilder接口负责组装 learner/actor/回放learning.pyMultiBCLearner与AquademLearner两阶段学习核心逻辑networks.pyFlaxEncoder网络多候选动作生成、DQN Q 网络actor.pyAquademActor执行离散选择 → 连续动作映射的动作生成utils.py环境创建、演示数据集加载与奖励稀疏化wrappers.pyAdroit 任务的稀疏奖励/成功奖励包装器requirements.txtPython 依赖清单整个算法可以概括为Encoder 预训练阶段 离散 RL 阶段两阶段流水线下面分别深入剖析。五、核心机制一多分类行为克隆MultiBC学习动作候选5.1 Encoder 网络结构动作候选生成器定义在 networks.py 的EncoderFlaxlinen.Module中其结构为共享 torso num_actions个独立 head输入观测先过一个共享 torso默认torso_layer_sizes(256,)即一层 256 维全连接 ReLU随后并行搭建num_actions个 head默认head_layer_sizes(256,)每个 head 最终输出action_dim维连续动作最后通过jnp.stack(actions, axis-1)把所有候选动作堆叠为形状[batch, action_dim, num_actions]的张量。也就是说给定一个观测Encoder 一次给出num_actions个不同的候选动作每个候选对应一种可能的行为模式。网络在输入层与隐藏层均使用 dropout默认input_dropout_rate0.1、hidden_dropout_rate0.1以鼓励多个 head 分化、避免坍缩到同一动作。5.2 Softmin 距离损失MultiBCLearnerlearning.py使用一种软最小值softmin损失训练 Encoder。其核心函数aqualosslearning.py的计算过程是计算每个候选动作与演示动作的平方 L2 距离action_distances sum((predicted - action)^2)对num_actions个距离做 softmin 聚合softmin_action_distances temperature * ( jax.nn.logsumexp(-action_distances / temperature) - jnp.log(num_actions)) loss -softmin_action_distances直觉上损失鼓励至少有一个候选动作接近专家动作因为 softmin 近似于取最小值而非强迫所有 head 都预测同一个动作——这正是它能学到多模态演示分布的关键。参数temperature默认0.001控制 softmin 的锐利程度。5.3 预训练阶段参数预训练由AquademConfigconfig.py中的字段控制配置项默认值说明num_actions10学习到的候选动作数量量化粒度encoder_learning_rate3e-4MultiBC 使用的 Adam 学习率encoder_batch_size256预训练 batch 大小encoder_num_steps50_000MultiBC 预训练总步数encoder_eval_every1_000预训练内部记录频率每步包含encoder_eval_every次 SGDtemperature0.001softmin 聚合的温度在 builder.py 中Encoder 使用optax.adam(encoder_learning_rate)优化learning.py 中AquademLearner构造时即先完成整个预训练以encoder_batch_size * encoder_eval_every条演示构造数据集循环encoder_num_steps // encoder_eval_every次调用MultiBCLearner.step()每次内部通过jax.jit编译执行encoder_eval_every次 SGD见process_multiple_batcheslearning.py。六、核心机制二离散 RL 在候选动作上学习6.1 从连续演示到离散标签预训练完成后AQuaDem 需要把回放数据翻译成离散 RL 可用的形式。_generate_aquadem_sampleslearning.py以概率demonstration_ratio默认0.25从演示数据中采样并将连续专家动作映射为距离最近的候选动作索引discrete_actions np.argmin( np.linalg.norm(continuous_actions_candidates - demonstrations.action[:, :, None], axis1), axis-1)同时若配置了min_demo_reward还会把演示样本的奖励下限提升到该值reward max(min_demo_reward, reward)从而鼓励 agent 加入专家的支撑集。其余情况下直接透传回放缓冲区中的交互样本。6.2 离散 RLMunchausen Q-Learning在离散化后的动作空间上AQuaDem 复用了 Acme 的 DQN 实现run_aquadqn.pyloss_fn dqn.losses.MunchausenQLearning(max_abs_reward100.) dqn_config dqn.DQNConfig( min_replay_size1000, n_step3, num_sgd_steps_per_step8, learning_rate1e-4, samples_per_insert256) rl_agent dqn.DQNBuilder(configdqn_config, loss_fnloss_fn)其中关键的超参数含义为参数默认值说明min_replay_size1000回放缓冲区最少积累多少样本后开始学习n_step33 步回报n-step TDnum_sgd_steps_per_step8每个环境步执行的 SGD 次数learning_rate1e-4DQN 学习率samples_per_insert256每次插入回放时采样的样本数max_abs_reward100.Munchausen 损失的奖励裁剪上限Q 网络由make_q_networknetworks.py构建默认使用LayerNormMLP隐藏层为(512, 512, 256)输出维度等于num_actions。值得注意的是源码保留了architectureMLP分支注释为 AQuaOff architecture即论文后续工作 AQuaOff 的变体可通过该参数切换网络架构。6.3 两阶段的数据流AquademLearnerlearning.py把两阶段串成完整闭环构造MultiBCLearner并完成encoder_num_steps步预训练见 5.3将演示迭代器与回放迭代器喂给_generate_aquadem_samples生成离散标签 混合演示/交互的学习数据流lfd_iteratorLearning from Demonstrations用lfd_iterator作为离散 RL learner 的数据源此后每次step()只推进离散 RL 的学习。因此num_actions、demonstration_ratio、min_demo_reward是决定演示在离散 RL 阶段发挥多大作用的三个核心旋钮。七、Actor如何把离散决策变成连续动作训练与评估时执行动作的组件是AquademActoractor.py它包装了一个离散动作 actorselect_action(observation)先调用内部离散策略DQN 的default_behavior_policy选出候选索引discrete_action再用 Encoder 对观测输出全部候选动作取第discrete_action个作为最终连续动作def aquadem_policy(params, observation, discrete_action): predicted_actions networks.encoder.apply(params, observation) return predicted_actions[..., discrete_action]环境交互回放时记录的是离散动作observe()使用self._last_discrete_action从而保证离散 RL 的训练数据一致。训练期间离散策略使用exploration_epsilon 0.01的 ε-greedy 探索run_aquadqn.py评估时 ε 设为 0 执行纯贪心策略。AquademActor中的 encoder 变量通过VariableClient从 learner 同步且由于 Encoder 预训练完成后不再更新update_period被设为极大的值以永不更新builder.py。八、演示数据与环境处理8.1 D4RL Adroit 数据集加载演示数据通过 TFDS 加载 D4RL Adroit 数据集utils.py 的_d4rl_dataset_name把环境名转换为 TFDS 数据集名例如door-human-v1→d4rl_adroit_door/v1-humanhammer-human-v1→d4rl_adroit_hammer/v1-humanget_make_demonstrations_fnutils.py负责加载 TFDS 数据集 →可选截取前num_demonstrations条 →稀疏化奖励按任务阈值把稠密奖励转为 0/1→ 转换为 Acme Transition 迭代器最终返回一个按 batch 大小生成随机演示 batch 的函数。8.2 奖励稀疏化与任务阈值Adroit 灵巧手任务door、hammer、pen、relocate的稀疏奖励阈值定义在 utils.pySPARSE_REWARD_THRESHOLDS {door: 15, hammer: 50, pen: 30, relocate: 5}演示数据中的奖励被稀疏化为reward threshold的 0/1 指示。训练环境则通过AdroitSparseRewardWrapperwrappers.py直接用环境自身的info[goal_achieved]作为奖励与演示数据稀疏化阈值保持一致评估环境再叠加SuccessRewardWrapper把整条轨迹是否成功作为评估指标。8.3 环境包装链make_environmentutils.py构建环境的完整包装链为gym.make(task) → AdroitSparseRewardWrappergoal_achieved 作为奖励 → GymWrapper转为 dm_env 接口 → CanonicalSpecWrapperclipTrue裁剪动作到 spec 范围 → SinglePrecisionWrapper单精度 → 仅评估时SuccessRewardWrapper整条轨迹成功即 1九、配置参数速查与调参建议综合 config.py 与 run_aquadqn.pyAQuaDem 的全部核心可调参数汇总如下算法层AquademConfig在代码中修改num_actions候选动作数越大表达能力越强但离散 RL 难度越高、encoder_learning_rate、encoder_batch_size、encoder_num_steps预训练步数、temperaturesoftmin 温度、demonstration_ratio演示样本混合比例、min_demo_reward演示奖励下限。RL 层DQNConfigmin_replay_size、n_step、num_sgd_steps_per_step、learning_rate、samples_per_insert、max_abs_reward。命令行层--workdir、--env_name、--num_demonstrations、--num_steps、--eval_every、--seed。常见调参思路基于源码逻辑推断候选动作数num_actions越大Encoder 对多模态演示的覆盖越充分但离散 RL 的 action space 也越大可相应增加num_sgd_steps_per_step或总步数--num_steps演示充足时可提高demonstration_ratio默认 0.25并设置合理的min_demo_reward强化演示对离散 RL 的引导temperature影响 softmin 的锐利度过小可能导致训练不稳定过大会让损失退化为均值近似切换环境时务必确认env_name对应的任务在SPARSE_REWARD_THRESHOLDS中已有阈值当前仅支持 door、hammer、pen、relocate 四个 Adroit 任务。十、总结AQuaDem 通过多分类 BC 学习动作候选 离散 RL 学习候选选择的两阶段设计把连续控制问题转化为离散决策问题同时让专家演示同时作用于 Encoder 预训练与 RL 训练两个环节。仓库源码以 Acme/JAX 生态实现结构清晰、模块边界明确networks.py负责多候选动作生成与 Q 网络learning.py承载两阶段学习核心actor.py完成离散到连续的最终映射utils.py/wrappers.py解决 Adroit 演示数据与环境奖励的一致性问题。理解这些模块后你可以参照 run_aquadqn.py 的组装方式将 AQuaDem 适配到自定义的连续控制任务需自行准备对应的演示数据集与奖励阈值。赞分享人工智能深度学习NLP计算机视觉强化学习【免费下载链接】google-researchGoogle Research项目地址https://gitcode.com/gh_mirrors/go/google-research点击查看免费下载相关推荐突破动作空间限制CleanRL离散与连续动作算法实战指南突破动作空间限制CleanRL离散与连续动作算法实战指南 你是否还在为强化学习算法适配不同动作空间而头疼一文掌握CleanRL如何优雅处理离散与连续动作问题人工智能机器学习深度学习强化学习突破连续动作难题PPO算法实战指南突破连续动作难题PPO算法实战指南 你是否在强化学习任务中遇到过机械臂控制、自动驾驶等连续动作空间难题传统算法要么收敛速度慢要么策略不稳定。本文将带你掌握教程机器学习深度学习Easy-RL连续动作空间DDPG与TD3算法实战Easy RL连续动作空间DDPG与TD3算法实战 本文深入探讨了强化学习在连续动作空间中的挑战与解决方案。首先分析了连续动作空间面临的动作选择组合爆炸、价值人工智能强化学习深度学习教程创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考