面向强化学习的状态空间建模：RSSM的先容和PyTorch实现

循环状态空间模子（Recurrent State Space Models, RSSM）最初由 Danijar Hafer 等人在论文《Learning Latent Dynamics for Planning from Pixels》中提出。该模子在当代基于模子的强化学习（Model-Based Reinforcement Learning, MBRL）中发挥着关键作用，其告急目的是构建可靠的情况动态推测模子。通过这些学习得到的模子，智能体可以大概模拟未来轨迹并举行前瞻性的举动规划。

下面我们就来用一个实际案例来先容RSSM。
情况设置

情况设置是实现过程中的主要步调。我们这里用易于使用的 Gym API。为了进步实现服从，计划了多个模块化的包装器（wrapper），用于初始化参数并将观察效果调解为指定格式。
InitialWrapper 的计划答应在不实验任何动作的情况下举行特定命目的观察，同时支持在返回观察效果之前多次重复同一动作。这种计划对于相应具有显着延伸特性的情况特殊有用。
PreprocessFrame 包装器负责将观察效果转换为正确的数据范例（本文中使用 numpy 数组），并支持灰度转换功能。

1. classInitialWrapper(gym.Wrapper):  
2.      def__init__(self, env: gym.Env, no_ops: int=0, repeat: int=1):  
3.          super(InitialWrapper, self).__init__(env)  
4.          self.repeat=repeat  
5.          self.no_ops=no_ops  
7.          self.op_counter=0  
8.    
9.      defstep(self, action: ActType) ->Tuple[ObsType, float, bool, bool, dict]:  
10.          ifself.op_counter<self.no_ops:  
11.              obs, reward, done, info=self.env.step(0)  
12.              self.op_counter+=1  
13.    
14.          total_reward=0.0  
15.          done=False  
16.          for_inrange(self.repeat):  
17.              obs, reward, done, info=self.env.step(action)  
18.              total_reward+=reward  
19.              ifdone:  
20.                  break  
21.    
22.          returnobs, total_reward, done, info  
25. classPreprocessFrame(gym.ObservationWrapper):  
26.      def__init__(self, env: gym.Env, new_shape: Sequence[int] = (128, 128, 3), grayscale: bool=False):  
27.          super(PreprocessFrame, self).__init__(env)  
28.          self.shape=new_shape  
29.          self.observation_space=gym.spaces.Box(low=0.0, high=1.0, shape=self.shape, dtype=np.float32)  
30.          self.grayscale=grayscale  
31.    
32.          ifself.grayscale:  
33.              self.observation_space=gym.spaces.Box(low=0.0, high=1.0, shape=(*self.shape[:-1], 1), dtype=np.float32)  
34.    
35.      defobservation(self, obs: torch.Tensor) ->torch.Tensor:  
36.          obs=obs.astype(np.uint8)  
37.          new_frame=cv.resize(obs, self.shape[:-1], interpolation=cv.INTER_AREA)  
38.          ifself.grayscale:  
39.              new_frame=cv.cvtColor(new_frame, cv.COLOR_RGB2GRAY)  
40.              new_frame=np.expand_dims(new_frame, -1)  
41.    
42.          torch_frame=torch.from_numpy(new_frame).float()  
43.          torch_frame=torch_frame/255.0  
44.    
45.          returntorch_frame  
46.    
47. defmake_env(env_name: str, new_shape: Sequence[int] = (128, 128, 3), grayscale: bool=True, **kwargs):  
48.      env=gym.make(env_name, **kwargs)  
49.      env=PreprocessFrame(env, new_shape, grayscale=grayscale)  
50.      returnenv

复制代码

make_env 函数用于创建一个具有指定设置参数的情况实例。
模子架构

RSSM 的实现依赖于多个关键模子组件。详细来说，须要实现以下四个焦点模块：

• 原始观察编码器（Encoder）
  
• 动态模子（Dynamics Model）：通过确定性状态 h 和随机状态 s 对编码观察的时间依赖性举行建模
  
• 解码器（Decoder）：将随机状态和确定性状态映射回原始观察空间
  
• 嘉奖模子（Reward Model）：将随机状态和确定性状态映射到嘉奖值
  

RSSM 模子组件结构图。模子包罗随机状态 s 和确定性状态 h。
编码器实现

编码器接纳简朴的卷积神经网络（CNN）结构，将输入图像降维到一维嵌入体现。实现中使用了 BatchNorm 来提升训练稳固性。

1. classEncoderCNN(nn.Module):  
2.      def__init__(self, in_channels: int, embedding_dim: int=2048, input_shape: Tuple[int, int] = (128, 128)):  
3.          super(EncoderCNN, self).__init__()  
4.          # 定义卷积层结构
5.          self.conv1=nn.Conv2d(in_channels, 32, kernel_size=3, stride=2, padding=1)  
6.          self.conv2=nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)  
7.          self.conv3=nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)  
8.          self.conv4=nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)  
9.    
10.          self.fc1=nn.Linear(self._compute_conv_output((in_channels, input_shape[0], input_shape[1])), embedding_dim)  
11.    
12.          # 批标准化层
13.          self.bn1=nn.BatchNorm2d(32)  
14.          self.bn2=nn.BatchNorm2d(64)  
15.          self.bn3=nn.BatchNorm2d(128)  
16.          self.bn4=nn.BatchNorm2d(256)  
17.    
18.      def_compute_conv_output(self, shape: Tuple[int, int, int]):  
19.          withtorch.no_grad():  
20.              x=torch.randn(1, shape[0], shape[1], shape[2])  
21.              x=self.conv1(x)  
22.              x=self.conv2(x)  
23.              x=self.conv3(x)  
24.              x=self.conv4(x)  
25.    
26.              returnx.shape[1] *x.shape[2] *x.shape[3]  
28.      defforward(self, x):  
29.          x=torch.relu(self.conv1(x))  
30.          x=self.bn1(x)  
31.          x=torch.relu(self.conv2(x))  
32.          x=self.bn2(x)  
33.    
34.          x=torch.relu(self.conv3(x))  
35.          x=self.bn3(x)  
36.    
37.          x=self.conv4(x)  
38.          x=self.bn4(x)  
39.    
40.          x=x.view(x.size(0), -1)  
41.          x=self.fc1(x)  
42.    
43.          returnx

复制代码

解码器实现

解码器依照传统自编码器架构计划，其功能是将编码后的观察效果重修回原始观察空间。

1. classDecoderCNN(nn.Module):  
2.      def__init__(self, hidden_size: int, state_size: int,  embedding_size: int,  
3.                   use_bn: bool=True, output_shape: Tuple[int, int] = (3, 128, 128)):  
4.          super(DecoderCNN, self).__init__()  
5.    
6.          self.output_shape=output_shape  
7.    
8.          self.embedding_size=embedding_size  
9.          # 全连接层进行特征变换
10.          self.fc1=nn.Linear(hidden_size+state_size, embedding_size)  
11.          self.fc2=nn.Linear(embedding_size, 256* (output_shape[1] //16) * (output_shape[2] //16))  
12.    
13.          # 反卷积层进行上采样
14.          self.conv1=nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, output_padding=1)  # ×2  
15.          self.conv2=nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1)  # ×2  
16.          self.conv3=nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, output_padding=1)  # ×2  
17.          self.conv4=nn.ConvTranspose2d(32, output_shape[0], kernel_size=3, stride=2, padding=1, output_padding=1)  
18.    
19.          # 批标准化层
20.          self.bn1=nn.BatchNorm2d(128)  
21.          self.bn2=nn.BatchNorm2d(64)  
22.          self.bn3=nn.BatchNorm2d(32)  
23.    
24.          self.use_bn=use_bn  
26.      defforward(self, h: torch.Tensor, s: torch.Tensor):  
27.          x=torch.cat([h, s], dim=-1)  
28.          x=self.fc1(x)  
29.          x=torch.relu(x)  
30.          x=self.fc2(x)  
31.    
32.          x=x.view(-1, 256, self.output_shape[1] //16, self.output_shape[2] //16)  
33.    
34.          ifself.use_bn:  
35.              x=torch.relu(self.bn1(self.conv1(x)))  
36.              x=torch.relu(self.bn2(self.conv2(x)))  
37.              x=torch.relu(self.bn3(self.conv3(x)))  
38.    
39.          else:  
40.              x=torch.relu(self.conv1(x))  
41.              x=torch.relu(self.conv2(x))  
42.              x=torch.relu(self.conv3(x))  
43.    
44.          x=self.conv4(x)  
45.    
46.          returnx   

复制代码

嘉奖模子实现

嘉奖模子接纳了一个三层前馈神经网络结构，用于将随机状态 s 和确定性状态 h 映射到正态分布参数，进而通过采样得到嘉奖推测。

1. classRewardModel(nn.Module):  
2.      def__init__(self, hidden_dim: int, state_dim: int):  
3.          super(RewardModel, self).__init__()  
4.    
5.          self.fc1=nn.Linear(hidden_dim+state_dim, hidden_dim)  
6.          self.fc2=nn.Linear(hidden_dim, hidden_dim)  
7.          self.fc3=nn.Linear(hidden_dim, 2)  
8.    
9.      defforward(self, h: torch.Tensor, s: torch.Tensor):  
10.          x=torch.cat([h, s], dim=-1)  
11.          x=torch.relu(self.fc1(x))  
12.          x=torch.relu(self.fc2(x))  
13.          x=self.fc3(x)  
14.    
15.          returnx

复制代码

动态模子的实现

动态模子是 RSSM 架构中最复杂的组件，须要同时处理处罚先验和后验状态转移模子：

• 后验转移模子：在可以大概访问真实观察的情况下使用（告急在训练阶段），用于在给定观察和汗青状态的条件下近似随机状态的后验分布。
  
• 先验转移模子：用于近似先验状态分布，仅依赖于前一时候状态，不依赖于观察。这在无法获取后验观察的推理阶段使用。
  

这两个模子均通过单层前馈网络举行参数化，输出各自正态分布的均值和对数方差，用于状态 s 的采样。该实现接纳了简朴的网络结构，但可以根据须要扩展为更复杂的架构。
确定性状态接纳门控循环单位（GRU）实现。其输入包罗：

• 前一时候的隐蔽状态
  
• 独热编码动作
  
• 前一时候随机状态 s（根据是否可以获取观察来选择使用后验或先验状态）
  

这些输入信息足以让模子相识动作汗青和体系状态。以下是详细实当代码：

1. classDynamicsModel(nn.Module):  
2.      def__init__(self, hidden_dim: int, action_dim: int, state_dim: int, embedding_dim: int, rnn_layer: int=1):  
3.          super(DynamicsModel, self).__init__()  
4.    
5.          self.hidden_dim=hidden_dim  
6.            
7.          # 递归层实现，支持多层 GRU
8.          self.rnn=nn.ModuleList([nn.GRUCell(hidden_dim, hidden_dim) for_inrange(rnn_layer)])  
9.            
10.          # 状态动作投影层
11.          self.project_state_action=nn.Linear(action_dim+state_dim, hidden_dim)  
12.            
13.          # 先验网络：输出正态分布参数
14.          self.prior=nn.Linear(hidden_dim, state_dim*2)  
15.          self.project_hidden_action=nn.Linear(hidden_dim+action_dim, hidden_dim)  
16.            
17.          # 后验网络：输出正态分布参数
18.          self.posterior=nn.Linear(hidden_dim, state_dim*2)  
19.          self.project_hidden_obs=nn.Linear(hidden_dim+embedding_dim, hidden_dim)  
20.    
21.          self.state_dim=state_dim  
22.          self.act_fn=nn.ReLU()  
23.    
24.      defforward(self, prev_hidden: torch.Tensor, prev_state: torch.Tensor, actions: torch.Tensor,  
25.                  obs: torch.Tensor=None, dones: torch.Tensor=None):  
26.          """  
27.          动态模型的前向传播
28.          参数:  
29.              prev_hidden: RNN的前一隐藏状态，形状 (batch_size, hidden_dim)  
30.              prev_state: 前一随机状态，形状 (batch_size, state_dim)  
31.              actions: 独热编码动作序列，形状 (sequence_length, batch_size, action_dim)  
32.              obs: 编码器输出的观察嵌入，形状 (sequence_length, batch_size, embedding_dim)  
33.              dones: 终止状态标志
34.          """  
35.          B, T, _=actions.size()  # 用于无观察访问时的推理
36.    
37.          # 初始化存储列表
38.          hiddens_list= []  
39.          posterior_means_list= []  
40.          posterior_logvars_list= []  
41.          prior_means_list= []  
42.          prior_logvars_list= []  
43.          prior_states_list= []  
44.          posterior_states_list= []  
45.            
46.          # 存储初始状态
47.          hiddens_list.append(prev_hidden.unsqueeze(1))   
48.          prior_states_list.append(prev_state.unsqueeze(1))  
49.          posterior_states_list.append(prev_state.unsqueeze(1))  
50.    
51.          # 时序展开
52.          fortinrange(T-1):  
53.              # 提取当前时刻状态和动作
54.              action_t=actions[:, t, :]  
55.              obs_t=obs[:, t, :] ifobsisnotNoneelsetorch.zeros(B, self.embedding_dim, device=actions.device)  
56.              state_t=posterior_states_list[-1][:, 0, :] ifobsisnotNoneelseprior_states_list[-1][:, 0, :]  
57.              state_t=state_tifdonesisNoneelsestate_t* (1-dones[:, t, :])  
58.              hidden_t=hiddens_list[-1][:, 0, :]  
59.                
60.              # 状态动作组合
61.              state_action=torch.cat([state_t, action_t], dim=-1)  
62.              state_action=self.act_fn(self.project_state_action(state_action))  
63.    
64.              # RNN 状态更新
65.              foriinrange(len(self.rnn)):  
66.                  hidden_t=self.rnn[i](state_action, hidden_t)  
67.    
68.              # 先验分布计算
69.              hidden_action=torch.cat([hidden_t, action_t], dim=-1)  
70.              hidden_action=self.act_fn(self.project_hidden_action(hidden_action))  
71.              prior_params=self.prior(hidden_action)  
72.              prior_mean, prior_logvar=torch.chunk(prior_params, 2, dim=-1)  
73.    
74.              # 从先验分布采样
75.              prior_dist=torch.distributions.Normal(prior_mean, torch.exp(F.softplus(prior_logvar)))  
76.              prior_state_t=prior_dist.rsample
77. ()  
78.    
79.              # 后验分布计算
80.              ifobsisNone:  
81.                  posterior_mean=prior_mean  
82.                  posterior_logvar=prior_logvar  
83.              else:  
84.                  hidden_obs=torch.cat([hidden_t, obs_t], dim=-1)  
85.                  hidden_obs=self.act_fn(self.project_hidden_obs(hidden_obs))  
86.                  posterior_params=self.posterior(hidden_obs)  
87.                  posterior_mean, posterior_logvar=torch.chunk(posterior_params, 2, dim=-1)  
88.    
89.              # 从后验分布采样
90.              posterior_dist=torch.distributions.Normal(posterior_mean, torch.exp(F.softplus(posterior_logvar)))  
91.              posterior_state_t=posterior_dist.rsample
92. ()  
93.    
94.              # 保存状态
95.              posterior_means_list.append(posterior_mean.unsqueeze(1))  
96.              posterior_logvars_list.append(posterior_logvar.unsqueeze(1))  
97.              prior_means_list.append(prior_mean.unsqueeze(1))  
98.              prior_logvars_list.append(prior_logvar.unsqueeze(1))  
99.              prior_states_list.append(prior_state_t.unsqueeze(1))  
100.              posterior_states_list.append(posterior_state_t.unsqueeze(1))  
101.              hiddens_list.append(hidden_t.unsqueeze(1))  
102.    
103.          # 合并时序数据
104.          hiddens=torch.cat(hiddens_list, dim=1)  
105.          prior_states=torch.cat(prior_states_list, dim=1)  
106.          posterior_states=torch.cat(posterior_states_list, dim=1)  
107.          prior_means=torch.cat(prior_means_list, dim=1)  
108.          prior_logvars=torch.cat(prior_logvars_list, dim=1)  
109.          posterior_means=torch.cat(posterior_means_list, dim=1)  
110.          posterior_logvars=torch.cat(posterior_logvars_list, dim=1)  
111.    
112.          returnhiddens, prior_states, posterior_states, prior_means, prior_logvars, posterior_means, posterior_logvars

复制代码

须要特殊注意的是，这里的观察输入并非原始观察数据，而是颠末编码器处理处罚后的嵌入体现。这种计划可以大概有用低沉盘算复杂度并提升模子的泛化本领。
RSSM 团体架构

将前述组件整合为完备的 RSSM 模子。其焦点是

1. generate_rollout

复制代码

方法，负责调用动态模子并天生情况动态的埋伏体现序列。对于没有汗青埋伏状态的情况（通常发生在轨迹开始时），该方法会举行须要的初始化。下面是完备的实当代码：

1. classRSSM:       def__init__(self,                    encoder: EncoderCNN,                    decoder: DecoderCNN,                    reward_model: RewardModel,                    dynamics_model: nn.Module,                    hidden_dim: int,                    state_dim: int,                    action_dim: int,                    embedding_dim: int,                    device: str="mps"):           """           循环状态空间模子（RSSM）实现                  参数:             encoder: 确定性状态编码器             decoder: 观察重构解码器             reward_model: 嘉奖推测模子             dynamics_model: 状态动态模子             hidden_dim: RNN 隐蔽层维度             state_dim: 随机状态维度             action_dim: 动作空间维度             embedding_dim: 观察嵌入维度             device: 盘算装备         """           super(RSSM, self).__init__()              # 模子组件初始化         self.dynamics=dynamics_model           self.encoder=encoder           self.decoder=decoder           self.reward_model=reward_model              # 维度参数存储         self.hidden_dim=hidden_dim           self.state_dim=state_dim           self.action_dim=action_dim           self.embedding_dim=embedding_dim              # 模子迁徙至指定装备         self.dynamics.to(device)           self.encoder.to(device)           self.decoder.to(device)           self.reward_model.to(device)          defgenerate_rollout
2. (self, actions: torch.Tensor, hiddens: torch.Tensor=None, states: torch.Tensor=None,                            obs: torch.Tensor=None, dones: torch.Tensor=None):           """         天生状态序列睁开                  参数:             actions: 动作序列             hiddens: 初始隐蔽状态（可选）             states: 初始随机状态（可选）             obs: 观察序列（可选）             dones: 停止标志序列                      返回:             完备的状态睁开序列         """         # 状态初始化         ifhiddensisNone:               hiddens=torch.zeros(actions.size(0), self.hidden_dim).to(actions.device)              ifstatesisNone:               states=torch.zeros(actions.size(0), self.state_dim).to(actions.device)              # 执办法态模子睁开         dynamics_result=self.dynamics(hiddens, states, actions, obs, dones)           hiddens, prior_states, posterior_states, prior_means, prior_logvars, posterior_means, posterior_logvars=dynamics_result              returnhiddens, prior_states, posterior_states, prior_means, prior_logvars, posterior_means, posterior_logvars          deftrain(self):           """启用训练模式"""         self.dynamics.train()           self.encoder.train()           self.decoder.train()           self.reward_model.train()          defeval(self):           """启用评估模式"""         self.dynamics.eval()           self.encoder.eval()           self.decoder.eval()           self.reward_model.eval()          defencode(self, obs: torch.Tensor):           """观察编码"""         returnself.encoder(obs)          defdecode(self, state: torch.Tensor):           """状态解码为观察"""         returnself.decoder(state)          defpredict_reward(self, h: torch.Tensor, s: torch.Tensor):           """嘉奖推测"""         returnself.reward_model(h, s)          defparameters(self):           """返回全部可训练参数"""         returnlist(self.dynamics.parameters()) +list(self.encoder.parameters()) + \                list(self.decoder.parameters()) +list(self.reward_model.parameters())          defsave(self, path: str):           """模子状态生存"""         torch.save({               "dynamics": self.dynamics.state_dict(),               "encoder": self.encoder.state_dict(),               "decoder": self.decoder.state_dict(),               "reward_model": self.reward_model.state_dict()           }, path)          defload(self, path: str):           """模子状态加载"""         checkpoint=torch.load(path)           self.dynamics.load_state_dict(checkpoint["dynamics"])           self.encoder.load_state_dict(checkpoint["encoder"])           self.decoder.load_state_dict(checkpoint["decoder"])           self.reward_model.load_state_dict(checkpoint["reward_model"])

复制代码

这个实现提供了一个完备的 RSSM 框架，包罗了模子的训练、评估、状态生存和加载等根本功能。该框架可以作为根本结构，根据详细应用场景举行扩展和优化。
训练体系计划

RSSM 的训练体系告急包罗两个焦点组件：履历回放缓冲区（Experience Replay Buffer）和智能体（Agent）。此中，缓冲区负责存储汗青履历数据用于训练，而智能体则作为情况与 RSSM 之间的接口，实现数据网络计谋。
履历回放缓冲区实现

缓冲区接纳循环队列结构，用于存储和管理观察、动作、嘉奖和停止状态等数据。通过

1. sample

复制代码

方法可以随机采样训练序列。

1. classBuffer:       def__init__(self, buffer_size: int, obs_shape: tuple, action_shape: tuple, device: torch.device):           """         履历回放缓冲区初始化                  参数:             buffer_size: 缓冲区容量             obs_shape: 观察数据维度             action_shape: 动作数据维度             device: 盘算装备         """         self.buffer_size=buffer_size           self.obs_buffer=np.zeros((buffer_size, *obs_shape), dtype=np.float32)           self.action_buffer=np.zeros((buffer_size, *action_shape), dtype=np.int32)           self.reward_buffer=np.zeros((buffer_size, 1), dtype=np.float32)           self.done_buffer=np.zeros((buffer_size, 1), dtype=np.bool_)              self.device=device           self.idx=0          defadd(self, obs: torch.Tensor, action: int, reward: float, done: bool):           """         添加单步履历数据         """         self.obs_buffer[self.idx] =obs           self.action_buffer[self.idx] =action           self.reward_buffer[self.idx] =reward           self.done_buffer[self.idx] =done           self.idx= (self.idx+1) %self.buffer_size        defsample
2. (self, batch_size: int, sequence_length: int):           """         随机采样履历序列                  参数:             batch_size: 批量巨细             sequence_length: 序列长度                      返回:             履历数据元组 (observations, actions, rewards, dones)         """         # 随机选择序列起始位置         starting_idxs=np.random.randint(0, (self.idx%self.buffer_size) -sequence_length, (batch_size,))                    # 构建完备序列索引         index_tensor=np.stack([np.arange(start, start+sequence_length) forstartinstarting_idxs])                    # 提取数据序列         obs_sequence=self.obs_buffer[index_tensor]           action_sequence=self.action_buffer[index_tensor]           reward_sequence=self.reward_buffer[index_tensor]           done_sequence=self.done_buffer[index_tensor]              returnobs_sequence, action_sequence, reward_sequence, done_sequence        defsave(self, path: str):           """生存缓冲区数据"""         np.savez(path, obs_buffer=self.obs_buffer, action_buffer=self.action_buffer,                    reward_buffer=self.reward_buffer, done_buffer=self.done_buffer, idx=self.idx)          defload(self, path: str):           """加载缓冲区数据"""         data=np.load(path)           self.obs_buffer=data["obs_buffer"]           self.action_buffer=data["action_buffer"]           self.reward_buffer=data["reward_buffer"]           self.done_buffer=data["done_buffer"]           self.idx=data["idx"]

复制代码

智能体计划

智能体实现了数据网络和规划功能。当前实现接纳了简朴的随机计谋举行数据网络，但该框架支持扩展更复杂的计谋。

1. classPolicy(ABC):       """计谋基类"""     @abstractmethod       def__call__(self, obs):           pass      classRandomPolicy(Policy):       """随机采样计谋"""     def__init__(self, env: Env):           self.env=env          def__call__(self, obs):           returnself.env.action_space.sample
2. ()    classAgent:       def__init__(self, env: Env, rssm: RSSM, buffer_size: int=100000,                   collection_policy: str="random", device="mps"):           """         智能体初始化                  参数:             env: 情况实例             rssm: RSSM模子实例             buffer_size: 履历缓冲区巨细             collection_policy: 数据网络计谋范例             device: 盘算装备         """         self.env=env           # 计谋选择         matchcollection_policy:               case"random":                   self.rollout_policy=RandomPolicy(env)               case_:                   raiseValueError("Invalid rollout policy")              self.buffer=Buffer(buffer_size, env.observation_space.shape,                             env.action_space.shape, device=device)           self.rssm=rssm          defdata_collection_action(self, obs):           """实验数据网络动作"""         returnself.rollout_policy(obs)          defcollect_data(self, num_steps: int):           """         网络训练数据                  参数:             num_steps: 网络步数         """         obs=self.env.reset()           done=False              iterator=tqdm(range(num_steps), desc="Data Collection")           for_initerator:               action=self.data_collection_action(obs)               next_obs, reward, done, _, _=self.env.step(action)               self.buffer.add(next_obs, action, reward, done)               obs=next_obs               ifdone:                   obs=self.env.reset()          defimagine_rollout(self, prev_hidden: torch.Tensor, prev_state: torch.Tensor,                         actions: torch.Tensor):           """         实验想象睁开                  参数:             prev_hidden: 前一隐蔽状态             prev_state: 前一随机状态             actions: 动作序列                      返回:             完备的模子输出，包罗隐蔽状态、先验状态、后验状态等         """         hiddens, prior_states, posterior_states, prior_means, prior_logvars, \         posterior_means, posterior_logvars=self.rssm.generate_rollout
3. (             actions, prev_hidden, prev_state)              # 在想象阶段使用先验状态推测嘉奖         rewards=self.rssm.predict_reward(hiddens, prior_states)              returnhiddens, prior_states, posterior_states, prior_means, \                prior_logvars, posterior_means, posterior_logvars, rewards          defplan(self, num_steps: int, prev_hidden: torch.Tensor,               prev_state: torch.Tensor, actions: torch.Tensor):           """         实验规划                  参数:             num_steps: 规划步数             prev_hidden: 初始隐蔽状态             prev_state: 初始随机状态             actions: 动作序列                      返回:             规划得到的隐蔽状态和先验状态序列         """         hidden_states= []           prior_states= []              hiddens=prev_hidden           states=prev_state              for_inrange(num_steps):               hiddens, states, _, _, _, _, _, _=self.imagine_rollout(                 hiddens, states, actions)               hidden_states.append(hiddens)               prior_states.append(states)              hidden_states=torch.stack(hidden_states)           prior_states=torch.stack(prior_states)              returnhidden_states, prior_states

复制代码

这部分实现提供了完备的数据管理和智能体交互框架。通过履历回放缓冲区，可以高效地存储和重用汗青数据；通过智能体的抽象计谋接口，可以方便地扩展差异的数据网络计谋。同时智能体还实现了基于模子的想象睁开和规划功能，为后续的决定订定提供了根本。
训练器实现与实验

训练器计划

训练器是 RSSM 实现中的末了一个关键组件，负责调和模子训练过程。训练器吸收 RSSM 模子、智能体、优化器等组件，并实现详细的训练逻辑。

1. logging.basicConfig(       level=logging.INFO,       format="%(asctime)s - %(levelname)s - %(message)s",       handlers=[           logging.StreamHandler(),  # 控制台输出         logging.FileHandler("training.log", mode="w")  # 文件输出     ]   )      logger=logging.getLogger(__name__)    classTrainer:       def__init__(self, rssm: RSSM, agent: Agent, optimizer: torch.optim.Optimizer,                   device: torch.device):           """         训练器初始化                  参数:             rssm: RSSM 模子实例             agent: 智能体实例             optimizer: 优化器实例             device: 盘算装备         """         self.rssm=rssm           self.optimizer=optimizer           self.device=device           self.agent=agent           self.writer=SummaryWriter()  # tensorboard 日志记载器        deftrain_batch(self, batch_size: int, seq_len: int, iteration: int,                     save_images: bool=False):           """         单批次训练                  参数:             batch_size: 批量巨细             seq_len: 序列长度             iteration: 当前迭代次数             save_images: 是否生存重修图像         """         # 采样训练数据         obs, actions, rewards, dones=self.agent.buffer.sample
2. (batch_size, seq_len)              # 数据预处理处罚         actions=torch.tensor(actions).long().to(self.device)           actions=F.one_hot(actions, self.rssm.action_dim).float()           obs=torch.tensor(obs, requires_grad=True).float().to(self.device)           rewards=torch.tensor(rewards, requires_grad=True).float().to(self.device)           dones=torch.tensor(dones).float().to(self.device)              # 观察编码         encoded_obs=self.rssm.encoder(obs.reshape(-1, *obs.shape[2:]).permute(0, 3, 1, 2))           encoded_obs=encoded_obs.reshape(batch_size, seq_len, -1)              # 实验 RSSM 睁开         rollout=self.rssm.generate_rollout
3. (actions, obs=encoded_obs, dones=dones)           hiddens, prior_states, posterior_states, prior_means, prior_logvars, \         posterior_means, posterior_logvars=rollout              # 重构观察         hiddens_reshaped=hiddens.reshape(batch_size*seq_len, -1)           posterior_states_reshaped=posterior_states.reshape(batch_size*seq_len, -1)           decoded_obs=self.rssm.decoder(hiddens_reshaped, posterior_states_reshaped)           decoded_obs=decoded_obs.reshape(batch_size, seq_len, *obs.shape[-3:])              # 嘉奖推测         reward_params=self.rssm.reward_model(hiddens, posterior_states)           mean, logvar=torch.chunk(reward_params, 2, dim=-1)           logvar=F.softplus(logvar)           reward_dist=Normal(mean, torch.exp(logvar))           predicted_rewards=reward_dist.rsample
4. ()              # 可视化         ifsave_images:               batch_idx=np.random.randint(0, batch_size)               seq_idx=np.random.randint(0, seq_len-3)               fig=self._visualize(obs, decoded_obs, rewards, predicted_rewards,                                  batch_idx, seq_idx, iteration, grayscale=True)               ifnotos.path.exists("reconstructions"):                   os.makedirs("reconstructions")               fig.savefig(f"reconstructions/iteration_{iteration}.png")               self.writer.add_figure("Reconstructions", fig, iteration)               plt.close(fig)              # 盘算丧失         reconstruction_loss=self._reconstruction_loss(decoded_obs, obs)           kl_loss=self._kl_loss(prior_means, F.softplus(prior_logvars),                                 posterior_means, F.softplus(posterior_logvars))           reward_loss=self._reward_loss(rewards, predicted_rewards)              loss=reconstruction_loss+kl_loss+reward_loss              # 反向传播和优化         self.optimizer.zero_grad()           loss.backward()           nn.utils.clip_grad_norm_(self.rssm.parameters(), 1, norm_type=2)           self.optimizer.step()              returnloss.item(), reconstruction_loss.item(), kl_loss.item(), reward_loss.item()          deftrain(self, iterations: int, batch_size: int, seq_len: int):           """         实验完备训练过程                  参数:             iterations: 迭代总次数             batch_size: 批量巨细             seq_len: 序列长度         """         self.rssm.train()           iterator=tqdm(range(iterations), desc="Training", total=iterations)           losses= []           infos= []           last_loss=float("inf")                    foriiniterator:               # 实验单批次训练             loss, reconstruction_loss, kl_loss, reward_loss=self.train_batch(                 batch_size, seq_len, i, save_images=i%100==0)                  # 记载训练指标             self.writer.add_scalar("Loss", loss, i)               self.writer.add_scalar("Reconstruction Loss", reconstruction_loss, i)               self.writer.add_scalar("KL Loss", kl_loss, i)               self.writer.add_scalar("Reward Loss", reward_loss, i)                  # 生存最佳模子             ifloss<last_loss:                   self.rssm.save("rssm.pth")                   last_loss=loss                  # 记载详细信息             info= {                   "Loss": loss,                   "Reconstruction Loss": reconstruction_loss,                   "KL Loss": kl_loss,                   "Reward Loss": reward_loss               }               losses.append(loss)               infos.append(info)                  # 定期输出训练状态             ifi%10==0:                   logger.info("\n----------------------------")                   logger.info(f"Iteration: {i}")                   logger.info(f"Loss: {loss:.4f}")                   logger.info(f"Running average last 20 losses: {sum(losses[-20:]) /20: .4f}")                   logger.info(f"Reconstruction Loss: {reconstruction_loss:.4f}")                   logger.info(f"KL Loss: {kl_loss:.4f}")                   logger.info(f"Reward Loss: {reward_loss:.4f}")  ### 实验示例  以下是一个在CarRacing情况中训练RSSM的完备示例：  ```python # 情况初始化 env=make_env("CarRacing-v2", render_mode="rgb_array", continuous=False, grayscale=True)    # 模子参数设置 hidden_size=1024   embedding_dim=1024   state_dim=512    # 模子组件实例化 encoder=EncoderCNN(in_channels=1, embedding_dim=embedding_dim)   decoder=DecoderCNN(hidden_size=hidden_size, state_size=state_dim,                       embedding_size=embedding_dim, output_shape=(1,128,128))   reward_model=RewardModel(hidden_dim=hidden_size, state_dim=state_dim)   dynamics_model=DynamicsModel(hidden_dim=hidden_size, state_dim=state_dim,                                action_dim=5, embedding_dim=embedding_dim)    # RSSM 模子构建 rssm=RSSM(dynamics_model=dynamics_model,               encoder=encoder,               decoder=decoder,               reward_model=reward_model,               hidden_dim=hidden_size,               state_dim=state_dim,               action_dim=5,               embedding_dim=embedding_dim)    # 训练设置 optimizer=torch.optim.Adam(rssm.parameters(), lr=1e-3)   agent=Agent(env, rssm)   trainer=Trainer(rssm, agent, optimizer=optimizer, device="cuda")    # 数据网络和训练 trainer.collect_data(20000)  # 网络 20000 步履历数据 trainer.save_buffer("buffer.npz")  # 生存履历缓冲区 trainer.train(10000, 32, 20)  # 实验 10000 次迭代训练

复制代码

总结

本文详细先容了基于 PyTorch 实现 RSSM 的完备过程。RSSM 的架构相比传统的 VAE 或 RNN 更为复杂，这告急源于其混淆了随机和确定性状态的特性。通过手动实现这一架构，我们可以深入明确其背后的理论根本及其强大之处。RSSM 可以大概递归地天生未来埋伏状态轨迹，这为智能体的举动规划提供了根本。
实现的长处在于其盘算负载适中，可以在单个斲丧级 GPU 上举行训练，在有富足时间的情况下乃至可以在 CPU 上运行。这一工作基于论文《Learning Latent Dynamics for Planning from Pixels》，该论文为 RSSM 类动态模子奠基了根本。后续的研究工作如《Dream to Control: Learning Behaviors by Latent Imagination》进一步发展了这一架构。这些改进的架构将在未来的研究中深入探究，由于它们对明确 MBRL 方法提供了告急的看法。
https://avoid.overfit.cn/post/8d8412f5ef6544e4ba097547a38830ac
作者：Lukas Bierling