训练配置（Training Configuration）

一句话定位

汇总世界模型训练中的关键超参数配置策略，包括学习率调度、KL权重管理、batch size选择和优化器配置，提供经验性默认值和调参建议。

前置依赖

1-复现笔记模板 - 配置记录格式
1-RSSM - RSSM训练目标理解
1-Dreamer系列 - 具体论文配置参考

核心思想

世界模型训练的挑战

世界模型训练涉及多个目标的平衡：

观察重建：确保潜在状态包含足够视觉信息
Dynamics学习：确保transition model准确
Reward预测：确保value function可靠
KL正则化：防止后验崩溃

核心矛盾：

KL weight 过高：压缩潜在表示，限制表示能力
KL weight 过低：先验-后验脱节，post. collapse
多个 loss 的 scale 差异大：训练不稳定

配置设计原则

从论文中提取默认值，但理解其敏感性：

参数类别	典型范围	调参优先级
学习率	1e-5 ~ 1e-3	高
KL weight	0.01 ~ 0.5	高
Batch size	16 ~ 256	中
Sequence length	16 ~ 128	中

数学推导

KL权重调度

固定KL权重的局限性：

固定在训练早期可能导致 posterior collapse，晚期可能限制表示学习。

KLannealing（KL退火）：

线性预热到目标值。

KL砍断（KL clipping）：

防止极端KL值影响训练。

损失函数归一化

Observation loss 归一化：

是观察维度，是序列长度。

Reward loss 归一化：

使用 reward variance 归一化。

优化器配置

Adam vs AdamW

标准Adam：

AdamW（带weight decay）：

世界模型推荐使用AdamW：

潜在状态的L2正则化有助于表示学习
防止权重膨胀导致的数值不稳定

学习率配置

经验性学习率范围：

组件	推荐范围	说明
视觉编码器	1e-6 ~ 5e-5	较小避免破坏预训练权重
Dynamics模型	5e-5 ~ 1e-4	中等
Reward预测器	5e-5 ~ 1e-4	中等
策略梯度	1e-6 ~ 1e-5	较小避免策略崩溃

分层学习率（Layer-wise LR Decay）：

其中。

学习率调度

Cosine Annealing：

适用于长期训练（>500k steps）。

Warmup + Cosine：

KL权重配置

典型配置方案

方案A：固定KL weight

kl_weight: 0.1
kl_free_nats: 0.0

适用于：训练稳定、数据充足

方案B：KL退火

kl_weight_start: 0.001
kl_weight_end: 0.5
kl_anneal_steps: 100000

适用于：避免 posterior collapse

方案C：KL砍断

kl_free_nats: 0.1
kl_clip: 1.0

适用于：控制 KL 量级

KL weight 敏感性分析

KL weight	效果	适用场景
< 0.01	潜在表示过分散，post. collapse	不推荐
0.05 ~ 0.1	平衡表示学习与重建	DM Control
0.1 ~ 0.5	强正则化，更结构化表示	复杂任务
> 0.5	表示过度压缩	不推荐

Batch Size配置

显存估算

影响因子：

模型参数量
Sequence length
观察维度（图像尺寸）
优化器状态（Adam: 4 bytes per param × 2 moments）

估算公式：

其中是参数量， batch size，序列长度，观察维度，动作维度。

经验值：

GPU	batch_size	seq_len	典型任务
RTX 3090	32-64	32	DMC 84×84
A100	128-256	64	DMC 84×84
V100	64-128	32	Atari 84×84

梯度累积

当显存不足时使用梯度累积：

batch_size: 8
grad_accum_steps: 4
effective_batch_size: 32

完整配置示例

DreamerV3风格配置

# world_model_training.yaml
model:
  name: dreamerv3
  stochastic_channels: 32
  deterministic_hidden: 512
  action_channels: 256
 
  encoder:
    cnn_channels: [48, 96, 192, 192]
    kernel_size: 4
    stride: 2
 
  dynamics:
    hidden_dim: 512
    num_layers: 1
 
  decoder:
    deconv_channels: [192, 96, 48, 48]
    output_activation: sigmoid
 
training:
  # 优化器
  optimizer:
    type: AdamW
    lr: 1e-4
    weight_decay: 0.0001
    betas: [0.9, 0.999]
    eps: 1e-5
 
  # 学习率调度
  lr_schedule:
    type: cosine
    warmup_steps: 5000
    min_lr: 1e-6
 
  # Batch配置
  batch_size: 64
  sequence_length: 32
  num_envs: 16
 
  # 损失权重
  losses:
    kl_weight: 0.1
    kl_free_nats: 0.0
    kl_balance: 0.5  # posterior vs prior weight ratio
 
    obs_weight: 1.0
    reward_weight: 1.0
    cont_weight: 0.1
 
  # 正则化
  grad_clip: 100.0
  ema_decay: 0.999
 
  # 动作采样
  action_sample_type: gaussian
  action_std_schedule: linear  # 从1.0降到0.1
 
environment:
  env_name: DM Control Suite
  num_eval_envs: 10
  num_train_envs: 16
 
  obs_spec:
    render: rgb_array
    size: [84, 84]
 
  action_spec:
    type: continuous
    max_abs_value: 1.0
 
evaluation:
  interval: 10000
  num_episodes: 10
  eval_max_steps: 1000

Transformer世界模型配置

# transformer_world_model.yaml
model:
  name: transformer_wm
  hidden_dim: 512
  num_layers: 6
  num_heads: 8
  seq_len: 64
  dropout: 0.1
 
  # 编码器
  encoder:
    type: vit
    patch_size: 16
    image_size: 224
 
  # 动作嵌入
  action_embed:
    dim: 64
    type: linear
 
training:
  optimizer:
    type: AdamW
    lr: 1e-4
    weight_decay: 0.01
    betas: [0.9, 0.999]
 
  lr_schedule:
    type: cosine
    warmup_steps: 10000
    min_lr: 1e-6
 
  batch_size: 32
  sequence_length: 64
 
  losses:
    obs_weight: 1.0
    action_weight: 0.5
    kl_weight: 0.05

章节摘要

本章提供了世界模型训练配置的完整指南：

优化器配置：AdamW优先，分层学习率策略
学习率调度：Cosine annealing + warmup是经验最优
KL权重管理：三种方案（固定/退火/砍断）及适用场景
Batch size选择：显存估算与梯度累积策略
完整配置示例：DreamerV3和TransformerWM的参考配置

正确的超参数配置是训练稳定收敛的关键，建议从标准配置开始，根据实验结果微调。

关键词

训练配置 超参数 AdamW 学习率调度 KL权重 KL退火 Batch Size 梯度累积 分层学习率 损失权重

Huajishe Notes

Explorer

2-训练配置