PPO 第三版：多环境

这次把 PPO 从“单个 Pendulum 环境一局一更新”的版本变成多环境并行，强化学习怎能不并行呢

1. 上一版的问题

训练主循环大概是这样的：

只创建一个 `Pendulum-v1` 环境。
每个 episode 跑最多 200 步。
每一步让 Actor 采样一个动作。
送进环境前，对动作做 `action.clip(ACTION_LOW, ACTION_HIGH)`。
一局结束后，用这一条轨迹做一次 PPO update。

这个版本已经比最初版稳定，因为它至少避免了 Pendulum 收到非法动作。但它还有几个明显问题。

首先采样效率低。PPO 是 on-policy 算法，每次更新都必须先采新数据。上一版只跑一个环境，一局最多 200 个 transition，样本量小，更新频率高，噪声也比较大。想让训练更稳，要么攒更多 episode 再更新，要么同时跑多个环境。

其次动作边界处理还不够完善和通用。上一版 Actor 里确实用了 tanh，但它作用在 mean 上，而且是手动乘 2：

1	mean = self.tanh(self.fc_mean(x)) * 2

这只能保证均值落在 Pendulum 的动作范围附近，不能保证从 Normal(mean, std) 里采样出来的 action 一定合法。只要 std 不是 0，采样动作就仍然可能超过 [-2, 2]。所以训练主循环仍然需要在环境外做 clip。

还有就是“策略采样动作”和“环境执行动作”不是同一套变换逻辑。上一版环境实际执行的是 clipped action，但策略分布本身还是按未 squash 的高斯动作来建模。虽然代码里把原始 action 存进 buffer，避免了 ratio 直接对不上，但动作边界这个变换发生在 Actor 外面，PPO 的 log_prob 并没有显式建模这个边界变换。

最后是没放GPU里跑

2. 多环境并行：一次 rollout 从 200 条样本变成 3200 条样本

train.py 最核心的变化，是把单环境：

1	env = gym.make(scenario)

改成了向量环境：

1 2	NUM_ENVS = 16 envs = gym.make_vec(scenario, num_envs=NUM_ENVS, vectorization_mode="sync")

现在每次迭代不是跑一个 Pendulum，而是 16 个 Pendulum 同时跑。ROLLOUT_STEPS = 200，所以一次 update 能拿到：

1	200 * 16 = 3200

条 transition。

让单次更新的数据量更大，以前一局只有 200 步，现在一轮就是 3200 步。PPO 的 advantage、ratio、loss 都建立在更大的样本集合上，梯度估计会更平滑。

3. 数据形状：从 `[T]` 轨迹到 `[T, N]` rollout

多环境以后，buffer 里存的数据不再是一条时间序列，而是一批并行时间序列，这个就和isaaclab中机器人训练类似了。

上一版里，reward、done、value 大致是：

1
2
3

rewards: [T]
values:  [T]
dones:   [T]

这一版变成：

rewards: [T, N]
values:  [T, N]
dones:   [T, N]
states:  [T, N, S]
actions: [T, N, A]

其中：

T = rollout 步数，也就是 200
N = 并行环境数，也就是 16
S = 状态维度
A = 动作维度

PPOAgent.get_action() 也因此从“给单个 state 出一个 action”，变成了“给一批 state 出一批 action”：

1
2
3

state = torch.as_tensor(state, dtype=torch.float32, device=self.device)
action, log_prob = self.actor.get_action(state)
value = self.critic(state)

返回值也保持批量形状：

1
2
3

action:   [N, A]
value:    [N]
log_prob: [N]

Actor 和 Critic 一次 forward 处理整个 batch。这才是后面 GPU 能发挥作用的基础。

4. GAE 向量化：一次反推 16 条轨迹

同样 GAE也变成多环境：

上一版

1 2	delta_t = r_t + gamma * V(s_{t+1}) - V(s_t) GAE_t = delta_t + gamma * lambda * GAE_{t+1}

这一版多环境后，last_gae 变成了长度为 N 的向量：

1 2	memo_advantages = torch.zeros((T, N), dtype=torch.float32, device=self.device) last_gae = torch.zeros(N, dtype=torch.float32, device=self.device)

反向循环仍然沿着时间维度走：

1 2	for t in reversed(range(T)): ...

但每一个时间步内部，delta、next_value、next_non_terminal 都是 [N]，也就是 16 个环境一起算。

这一步的意义是：时间依赖仍然保留，环境维度被向量化。

算 GAE 时，时间顺序不能打乱；但算完 GAE 后，每个 transition 已经有了自己的 advantage 和 return，就可以把 [T, N] 展平成一个大 batch：

memo_states        = memo_states.reshape(T * N, -1)
memo_actions       = memo_actions.reshape(T * N, -1)
memo_old_log_probs = memo_old_log_probs.reshape(-1)
memo_advantages    = memo_advantages.reshape(-1)
memo_returns       = memo_returns.reshape(-1)

这时训练样本数就是 T * N = 3200，再用 torch.randperm 打乱，切成 minibatch 更新。

5. GPU 化

因为之前的代码是在公司里电脑上写的训的，电脑只有CPU，所以在数据传递上就没有管太多，这次也完善一下，减少一些数据搬运。全放GPU里，但是这个环境是在CPU上跑的，实际上物理模拟还是无法加速，但是训练加速确实是存在的。

上一版中，很多中间数据先用 NumPy 整理，之后再转成 Tensor：

1
2
3

memo_states = np.array(...)
...
memo_states = torch.FloatTensor(memo_states).to(self.device)

这一版直接把 rollout 数据转成目标设备上的 Tensor：

1 2	memo_states = torch.as_tensor(np.array(memo_states, dtype=np.float32), device=self.device) memo_rewards = torch.as_tensor(np.array(memo_rewards, dtype=np.float32), device=self.device)

之后的 GAE、advantage 标准化、随机索引、minibatch 切片、Actor/Critic forward、loss 和反向传播，都在同一个 device 上完成：

1	indices = torch.randperm(batch_total, device=self.device)

6. 动作 squash ：去掉手动动作clip

上一版 Actor 已经有一个 tanh：

1 2	mean = self.tanh(self.fc_mean(x)) * 2 std = self.softplus(self.fc_std(x)) + 1e-5

这个写法的含义是：网络先输出均值，再把均值压到 [-2, 2]。它比完全无约束的均值更稳定，但还不是完整的 squashed policy。

原因在于，后面采样时用的是：

1	action = dist.sample()

Normal(mean, std) 的采样结果不是只会落在 [-2, 2] 里。std 只要不是 0，就仍然有概率采到边界外的动作。所以上一版送环境前还要：

1	clipped_action = action.clip(ACTION_LOW, ACTION_HIGH)

这一版不是简单“新增 tanh”，而是把 tanh 从“只约束 mean”推进到“约束最终采样动作”。具体做法是：先从无界正态分布采样 raw_action，再经过 tanh 压到 [-1, 1]，最后线性映射到环境动作范围：

1 2	raw_action = dist.sample() action = self._squash_action(raw_action)

对应实现是：

1 2	def _squash_action(self, raw_action): return torch.tanh(raw_action) * self.action_scale + self.action_bias

action_scale 和 action_bias 来自动作空间：

1 2	self.register_buffer("action_scale", (action_high - action_low) / 2.0, persistent=False) self.register_buffer("action_bias", (action_high + action_low) / 2.0, persistent=False)

对 Pendulum 来说，动作范围是 [-2, 2]，所以 tanh(raw_action) 会被映射到 [-2, 2]。

这样做以后，train.py 里不再需要手动 clip：

1	next_state, reward, terminated, truncated, _ = envs.step(action)

这比外部 clip 更自然，因为策略分布本身就产生合法动作。环境收到的 action 和 buffer 里存的 action 是同一个 action，行为策略和训练数据之间更一致。

所以这一节更准确的表述是：上一版有 tanh(mean) * 2 这种手动均值约束；这一版完善成了 Actor 内部对采样动作做 tanh squash，并用 action_scale/action_bias 适配环境动作上下界。

7. log_prob 修正：tanh 变换后不能直接用原始正态概率

只做 tanh squash 还不够。因为 PPO 的 ratio 依赖动作的 log probability：

1	ratio = torch.exp(new_log_probs - memo_old_log_probs[batch])

如果动作经过了 tanh 变换，动作空间的概率密度也被变换了。不能再直接用：

1	dist.log_prob(action)

这一版新增了 _squashed_log_prob()：

def _squashed_log_prob(self, dist, raw_action):
    squashed = torch.tanh(raw_action)
    log_prob = dist.log_prob(raw_action)
    log_det_jacobian = torch.log(self.action_scale * (1.0 - squashed.pow(2)) + self.EPS)
    return (log_prob - log_det_jacobian).sum(dim=-1)

这里的核心是 change-of-variables：

1	log pi(a) = log pi(u) - log \| da / du \|

其中 u 是无界的 raw_action，a 是 squash 后的环境动作。
这一步让 PPO ratio 比较的是“同一个动作在 squashed 策略下的概率”，而不是拿变换前后的两个空间硬比。

更新阶段还有一个配套函数：

def evaluate_actions(self, states, actions):
    mean, std = self.forward(states)
    dist = Normal(mean, std)
    raw_actions = self._inverse_squash_action(actions)
    log_probs = self._squashed_log_prob(dist, raw_actions)
    entropy = dist.entropy().sum(dim=-1).mean()
    return log_probs, entropy