代码架构建议（Code Architecture Guidelines）
一句话定位
提供世界模型项目的模块化代码架构设计，包含world model模块、policy模块、training loop的组织和接口定义，确保代码可复用性和工程可维护性。
前置依赖
核心思想
代码架构原则
模块化：
- 各组件（encoder、dynamics、decoder）独立
- 接口清晰，依赖最小化
- 便于单元测试和替换
可配置：
- 所有超参数通过config管理
- 不硬编码任何数值
- 配置即文档
可复现：
- 固定随机种子
- 完整记录依赖版本
- 环境可重建
三层架构
世界模型项目
├── Model Layer（世界模型本身）
│ ├── Encoder（观察编码）
│ ├── Dynamics（状态转移）
│ ├── Posterior / Prior（潜在分布）
│ └── Decoder（观察/奖励解码）
│
├── Policy Layer（策略/动作选择）
│ ├── Actor（策略网络）
│ └── Critic（价值函数）
│
└── Training Layer（训练逻辑）
├── DataLoader（数据管道）
├── TrainLoop（训练循环）
└── Evaluation（评估）
模块化设计
WorldModel Module
# world_model.py
class WorldModel(nn.Module):
"""世界模型主模块，组合所有子组件"""
def __init__(self, config):
super().__init__()
self.config = config
# 观察编码器
self.encoder = Encoder(
obs_dim=config.obs_dim,
hidden_dim=config.hidden_dim,
num_layers=config.num_encoder_layers,
)
# 潜在动力学
self.dynamics = RSSMDynamics(
stochastic_dim=config.stochastic_dim,
deterministic_dim=config.deterministic_dim,
action_dim=config.action_dim,
)
# 解码器
self.decoder = Decoder(
hidden_dim=config.hidden_dim,
obs_dim=config.obs_dim,
)
# Reward预测器
self.reward_predictor = RewardPredictor(
hidden_dim=config.hidden_dim,
)
def forward(self, obs_seq, action_seq):
"""
前向传播
Args:
obs_seq: (B, T, obs_dim) 观察序列
action_seq: (B, T, action_dim) 动作序列
Returns:
predictions: dict 包含重建、reward等
"""
# 编码观察
enc_obs = self.encoder(obs_seq)
# 潜在状态序列
post_seq, prior_seq, state_seq = self.dynamics(enc_obs, action_seq)
# 重建观察
recon_obs = self.decoder(state_seq)
# 预测reward
pred_reward = self.reward_predictor(state_seq)
return {
"posterior": post_seq,
"prior": prior_seq,
"states": state_seq,
"recon_obs": recon_obs,
"pred_reward": pred_reward,
}
def imagine(self, state, action_seq):
"""
想象 rollout（不需要真实观察）
Args:
state: 初始状态
action_seq: 动作序列
Returns:
imagined_obs: 想象中的观察
imagined_reward: 想象中的reward
"""
# 使用 prior 进行想象
return self.dynamics.imagine(state, action_seq)RSSMDynamics Module
# dynamics/rssm.py
class RSSMDynamics(nn.Module):
"""RSSM潜在动力学模块"""
def __init__(self, stochastic_dim, deterministic_dim, action_dim):
super().__init__()
self.stochastic_dim = stochastic_dim
self.deterministic_dim = deterministic_dim
# 确定性模型：LSTM或Transformer
self.rnn = nn.LSTM(
input_size=stochastic_dim + action_dim,
hidden_size=deterministic_dim,
num_layers=1,
batch_first=True,
)
# 后验网络：给定观察输出潜在分布
self.posterior_net = nn.Sequential(
nn.Linear(deterministic_dim + obs_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, stochastic_dim * 2), # mean + std
)
# 先验网络：仅基于状态输出潜在分布
self.prior_net = nn.Sequential(
nn.Linear(deterministic_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, stochastic_dim * 2),
)
def forward(self, obs_enc, action_seq):
"""
前向传播
Returns:
posterior_seq, prior_seq, hidden_seq
"""
T = obs_enc.shape[1]
hiddens = []
posteriors = []
priors = []
hidden = torch.zeros(1, obs_enc.shape[0], self.deterministic_dim)
for t in range(T):
# 混合分布：从 posterior 获取 z_t
if t == 0:
# 第一步没有观察，使用先验
z_t = torch.zeros_like(action_seq[:, 0])
else:
# 训练时用 posterior，推理时用先验
z_t = posteriors[t-1].rsample()
# 传入RNN
rnn_input = torch.cat([z_t, action_seq[:, t]], dim=-1)
hidden = self.rnn(rnn_input.unsqueeze(1), hidden)[0]
# 先验分布（基于 hidden）
prior_logits = self.prior_net(hidden)
prior_dist = torch.distributions.Normal(*self._split_logits(prior_logits))
# 训练时用真实观察计算后验
if obs_enc is not None:
posterior_logits = self.posterior_net(
torch.cat([hidden, obs_enc[:, t]], dim=-1)
)
posterior_dist = torch.distributions.Normal(
*self._split_logits(posterior_logits)
)
else:
posterior_dist = prior_dist
posteriors.append(posterior_dist)
priors.append(prior_dist)
hiddens.append(hidden)
return (
torch.stack(posteriors),
torch.stack(priors),
torch.stack(hiddens),
)
def _split_logits(self, logits):
mean, std = logits.chunk(2, dim=-1)
std = F.softplus(std) + 1e-6
return mean, stdPolicy Module
# policy.py
class Policy(nn.Module):
"""策略网络"""
def __init__(self, state_dim, action_dim, hidden_dim):
super().__init__()
self.actor = nn.Sequential(
nn.Linear(state_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, action_dim),
nn.Tanh(), # 连续动作
)
def forward(self, state):
"""给定状态，输出动作"""
return self.actor(state)
class ValueNetwork(nn.Module):
"""价值网络（用于Critic）"""
def __init__(self, state_dim, hidden_dim):
super().__init__()
self.critic = nn.Sequential(
nn.Linear(state_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, 1),
)
def forward(self, state):
return self.critic(state)Training Loop
核心训练循环
# train.py
class WorldModelTrainer:
def __init__(self, config):
self.config = config
self.world_model = WorldModel(config).to(device)
self.policy = Policy(config).to(device)
# 优化器
self.optimizer = AdamW([
{"params": self.world_model.parameters(), "lr": config.wm_lr},
{"params": self.policy.parameters(), "lr": config.policy_lr},
], weight_decay=config.weight_decay)
# 数据加载器
self.dataset = ReplayBuffer(config)
self.dataloader = DataLoader(
self.dataset,
batch_size=config.batch_size,
shuffle=True,
num_workers=4,
)
# 日志
self.writer = SummaryWriter(config.log_dir)
self.global_step = 0
def train_step(self, batch):
"""单步训练"""
obs_seq = batch["obs"].to(device)
action_seq = batch["action"].to(device)
reward_seq = batch["reward"].to(device)
# 世界模型前向
wm_output = self.world_model(obs_seq, action_seq)
# 计算损失
loss_obs = self.compute_obs_loss(wm_output, obs_seq)
loss_reward = self.compute_reward_loss(wm_output, reward_seq)
loss_kl = self.compute_kl_loss(wm_output)
loss_total = (
self.config.obs_weight * loss_obs +
self.config.reward_weight * loss_reward +
self.config.kl_weight * loss_kl
)
# 反向传播
self.optimizer.zero_grad()
loss_total.backward()
nn.utils.clip_grad_norm_(self.world_model.parameters(), self.config.grad_clip)
self.optimizer.step()
# 记录日志
self.log_metrics({
"train/loss_total": loss_total,
"train/loss_obs": loss_obs,
"train/loss_reward": loss_reward,
"train/loss_kl": loss_kl,
})
return loss_total
def compute_kl_loss(self, wm_output):
"""KL损失，支持退火"""
posterior = wm_output["posterior"]
prior = wm_output["prior"]
kl_per_dim = torch.distributions.kl_divergence(posterior, prior).mean(-1)
kl_per_dim = kl_per_dim.clamp(min=self.config.kl_free_nats)
# 退火
kl_weight = self.get_kl_weight()
return (kl_per_dim * kl_weight).mean()
def get_kl_weight(self):
"""动态KL权重"""
if self.config.kl_anneal:
progress = min(1.0, self.global_step / self.config.kl_anneal_steps)
return self.config.kl_weight * progress
return self.config.kl_weight
def train(self):
"""完整训练流程"""
while self.global_step < self.config.total_steps:
for batch in self.dataloader:
loss = self.train_step(batch)
self.global_step += 1
# Eval
if self.global_step % self.config.eval_interval == 0:
self.evaluate()
# Save
if self.global_step % self.config.save_interval == 0:
self.save_checkpoint()
def evaluate(self):
"""评估"""
eval_env = self.config.eval_env
returns = []
for _ in range(self.config.num_eval_episodes):
obs, _ = eval_env.reset()
hidden = self.world_model.initial_state()
done = False
episode_return = 0
while not done:
action = self.policy(hidden)
obs, reward, done, _, _ = eval_env.step(action)
episode_return += reward
returns.append(episode_return)
self.log_metrics({
"eval/return_mean": np.mean(returns),
"eval/return_std": np.std(returns),
})
def log_metrics(self, metrics):
"""统一日志接口"""
for name, value in metrics.items():
self.writer.add_scalar(name, value, self.global_step)
def save_checkpoint(self):
"""保存checkpoint"""
torch.save({
"step": self.global_step,
"world_model": self.world_model.state_dict(),
"policy": self.policy.state_dict(),
"optimizer": self.optimizer.state_dict(),
}, f"{self.config.checkpoint_dir}/step_{self.global_step}.pt")配置文件结构
主配置文件
# config.yaml
project:
name: world_model_project
seed: 42
device: cuda:0
model:
stochastic_dim: 32
deterministic_dim: 512
action_dim: 6
encoder:
type: conv
channels: [48, 96, 192]
kernel_size: 4
stride: 2
dynamics:
type: rssm
rnn_type: lstm
decoder:
type: deconv
channels: [192, 96, 48, 48]
training:
total_steps: 1000000
batch_size: 64
sequence_length: 32
wm_lr: 1e-4
policy_lr: 1e-5
weight_decay: 0.0001
grad_clip: 100.0
kl_weight: 0.1
kl_free_nats: 0.0
kl_anneal: true
kl_anneal_steps: 100000
obs_weight: 1.0
reward_weight: 0.1
# 优化器
optimizer: adamw
betas: [0.9, 0.999]
eps: 1e-5
# 学习率调度
lr_schedule:
type: cosine
warmup_steps: 5000
min_lr: 1e-6
environment:
name: DM Control Suite/walker_walk
num_train_envs: 16
num_eval_envs: 10
logging:
log_dir: logs/tensorboard
checkpoint_dir: checkpoints
eval_interval: 1000
save_interval: 10000
log_interval: 100项目目录结构
world_model_project/
├── src/
│ ├── __init__.py
│ ├── world_model.py # WorldModel主类
│ ├── encoder.py # 观察编码器
│ ├── decoder.py # 解码器
│ ├── dynamics/
│ │ ├── __init__.py
│ │ ├── rssm.py # RSSM实现
│ │ └── transformer_dyn.py # 可选的Transformer dynamics
│ ├── policy.py # 策略网络
│ ├── training.py # 训练逻辑
│ └── utils.py # 工具函数
│
├── configs/
│ ├── default.yaml # 默认配置
│ ├── dmc_walker.yaml # DMC walker任务
│ └── atari_breakout.yaml # Atari breakout任务
│
├── scripts/
│ ├── train.py # 训练脚本
│ ├── eval.py # 评估脚本
│ └── generate_videos.py # 生成回放视频
│
├── tests/
│ ├── test_encoder.py
│ ├── test_dynamics.py
│ └── test_training.py
│
├── requirements.txt
├── config.yaml
└── README.md
依赖管理
requirements.txt
torch>=2.0.0
dm-haiku>=0.0.10
optax>=0.1.5
numpy>=1.24.0
pyyaml>=6.0
tensorboard>=2.14.0
wandb>=0.16.0
gymnasium>=0.29.0
环境隔离
# 使用venv或conda
conda create -n world_model python=3.9
conda activate world_model
pip install -r requirements.txt章节摘要
本章提供了世界模型项目的完整代码架构设计：
- 三层架构：Model Layer、Policy Layer、Training Layer的职责划分
- WorldModel Module：核心模块的接口设计和实现
- RSSMDynamics Module：动力学模块的详细实现
- Policy Module：策略网络的独立设计
- Training Loop：完整的训练循环代码
- 配置管理：YAML配置结构和参数管理
- 目录结构：标准化的项目组织方式
良好的代码架构是科研效率的基础，建议在新项目开始时就按照上述结构组织代码。
关键词
代码架构 模块化设计 WorldModel RSSM 训练循环 配置文件 项目结构 软件工程 可复现性 最佳实践