世界模型与 Model-Based RL：从预测到自进化

# Model-Based RL 的核心循环
class ModelBasedAgent:
    def __init__(self):
        self.world_model = WorldModel()  # 环境模型
        self.policy = Policy()            # 决策策略
        self.planner = Planner()          # 规划器
    
    def train(self, env, steps):
        for step in range(steps):
            # 1. 收集真实经验
            obs, action, reward, next_obs = env.step(action)
            self.world_model.update(obs, action, next_obs, reward)
            
            # 2. 在世界模型中想象训练
            for _ in range(10):
                sim_obs = self.world_model.sample_state()
                sim_action = self.planner.plan(self.world_model, sim_obs)
                sim_next, sim_reward = self.world_model.predict(sim_obs, sim_action)
                self.policy.update(sim_obs, sim_action, sim_reward, sim_next)
            
            # 3. 选择真实动作
            action = self.planner.plan(self.world_model, obs)

# DreamerV3 风格的世界模型训练
import torch
import torch.nn as nn

class WorldModel(nn.Module):
    """潜空间世界模型"""
    def __init__(self, obs_dim, action_dim, latent_dim=64):
        super().__init__()
        # 编码器：观察 -> 潜状态
        self.encoder = nn.Sequential(
            nn.Linear(obs_dim, 256),
            nn.ReLU(),
            nn.Linear(256, latent_dim * 2)
        )
        # 动态模型：潜状态 + 动作 -> 下一个潜状态
        self.dynamics = nn.GRUCell(latent_dim + action_dim, latent_dim)
        # 奖励预测器
        self.reward_head = nn.Linear(latent_dim, 1)
        # 观察重建器
        self.decoder = nn.Linear(latent_dim, obs_dim)
    
    def imagine_trajectory(self, h, actions, horizon=10):
        """在世界模型中推演未来轨迹"""
        trajectories = []
        for t in range(horizon):
            h = self.dynamics(torch.cat([h, actions[t]], dim=-1), h)
            r = self.reward_head(h)
            trajectories.append({'state': h, 'reward': r})
        return trajectories

import numpy as np

class SimpleModelBasedAgent:
    """简化的 Model-Based RL Agent"""
    def __init__(self, state_dim, action_dim):
        self.W = np.zeros((state_dim, state_dim + action_dim))
        self.b = np.zeros(state_dim)
        self.R = np.zeros(state_dim + action_dim)
        
    def update_world_model(self, states, actions, next_states, rewards):
        """用收集的经验更新世界模型"""
        X = np.hstack([states, actions])
        self.W = np.linalg.lstsq(X, next_states, rcond=None)[0].T
        self.b = np.mean(next_states - X @ self.W.T, axis=0)
        self.R = np.linalg.lstsq(X, rewards, rcond=None)[0]
    
    def plan(self, current_state, horizon=5, n_samples=50):
        """在世界模型中规划：随机采样动作序列，选最优"""
        best_reward = -np.inf
        best_action = 0
        for _ in range(n_samples):
            s = current_state.copy()
            total_reward = 0
            actions = []
            for _ in range(horizon):
                a = np.random.randint(0, 4)
                actions.append(a)
                a_onehot = np.zeros(4)
                a_onehot[a] = 1
                x = np.hstack([s, a_onehot])
                s = x @ self.W.T + self.b
                total_reward += x @ self.R
            if total_reward > best_reward:
                best_reward = total_reward
                best_action = actions[0]
        return best_action