MaskablePPO训练时StableBaselines3-contrib产生NaN值问题求助
2026-7-2
MaskablePPO训练自定义Gym环境时NaN值问题的解决方案
问题描述
使用StableBaselines3-contrib的MaskablePPO训练自定义Gym环境时，神经网络输出NaN值导致训练报错。已尝试以下方案但无效：
- 将
np.float64改为np.float32 - 调整gamma参数
- 降低学习率
use_expln=True方案因MaskablePPO不支持无法使用
相关代码
class custom(gym.Env): ... env = custom() def mask_fn(env: gym.Env) -> List[bool]: return env.valid_action_mask() env = ActionMasker(env, mask_fn) model = MaskablePPO(MaskableActorCriticPolicy, env, gamma=0.001, verbose=0) checkpoint_callback = CheckpointCallback(save_freq=10000, save_path='logs', name_prefix='rl_model') model.learn(500000, callback=checkpoint_callback) model.save("JOM")
报错信息
ValueError Traceback (most recent call last) <ipython-input-9-abee064644f3> in <cell line: 3>() 1 checkpoint_callback = CheckpointCallback(save_freq=10000, save_path='logs', 2 name_prefix='rl_model') ----> 3 model.learn(500000, callback=checkpoint_callback) 4 model.save("JOM") 8 frames /usr/local/lib/python3.10/dist-packages/sb3_contrib/ppo_mask/ppo_mask.py in learn(self, total_timesteps, callback, log_interval, tb_log_name, reset_num_timesteps, use_masking, progress_bar) 545 self.logger.dump(step=self.num_timesteps) 546 --> 547 self.train() 548 549 callback.on_training_end() /usr/local/lib/python3.10/dist-packages/sb3_contrib/ppo_mask/ppo_mask.py in train(self) 410 actions = rollout_data.actions.long().flatten() 411 --> 412 values, log_prob, entropy = self.policy.evaluate_actions( 413 rollout_data.observations, 414 actions, /usr/local/lib/python3.10/dist-packages/sb3_contrib/common/maskable/policies.py in evaluate_actions(self, obs, actions, action_masks) 331 latent_vf = self.mlp_extractor.forward_critic(vf_features) 332 --> 333 distribution = self._get_action_dist_from_latent(latent_pi) 334 if action_masks is not None: 335 distribution.apply_masking(action_masks) /usr/local/lib/python3.10/dist-packages/sb3_contrib/common/maskable/policies.py in _get_action_dist_from_latent(self, latent_pi) 244 """ 245 action_logits = self.action_net(latent_pi) --> 246 return self.action_dist.proba_distribution(action_logits=action_logits) 247 248 def _predict( /usr/local/lib/python3.10/dist-packages/sb3_contrib/common/maskable/distributions.py in proba_distribution(self, action_logits) 192 reshaped_logits = action_logits.view(-1, sum(self.action_dims)) 193 --> 194 self.distributions = [ 195 MaskableCategorical(logits=split) for split in th.split(reshaped_logits, tuple(self.action_dims), dim=1) 196 ] /usr/local/lib/python3.10/dist-packages/sb3_contrib/common/maskable/distributions.py in <listcomp>(.0) 193 194 self.distributions = [ --> 195 MaskableCategorical(logits=split) for split in th.split(reshaped_logits, tuple(self.action_dims), dim=1) 196 ] 197 return self /usr/local/lib/python3.10/dist-packages/sb3_contrib/common/maskable/distributions.py in __init__(self, probs, logits, validate_args, masks) 40 ): 41 self.masks: Optional[th.Tensor] = None --> 42 super().__init__(probs, logits, validate_args) 43 self._original_logits = self.logits 44 self.apply_masking(masks) /usr/local/lib/python3.10/dist-packages/torch/distributions/categorical.py in __init__(self, probs, logits, validate_args) 68 self._param.size()[:-1] if self._param.ndimension() > 1 else torch.Size() 69 ) --> 70 super().__init__(batch_shape, validate_args=validate_args) 71 72 def expand(self, batch_shape, _instance=None): /usr/local/lib/python3.10/dist-packages/torch/distributions/distribution.py in __init__(self, batch_shape, event_shape, validate_args) 66 valid = constraint.check(value) 67 if not valid.all(): --> 68 raise ValueError( 69 f"Expected parameter {param} " 70 f"({type(value).__name__} of shape {tuple(value.shape)}) " ValueError: Expected parameter logits (Tensor of shape (64, 2)) of distribution MaskableCategorical(logits: torch.Size([64, 2])) to satisfy the constraint IndependentConstraint(Real(), 1), but found invalid values: tensor([[nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan], [nan, nan]], grad_fn=<SubBackward0>)
可行解决方案
1. 修正gamma参数（优先级最高）
你当前设置的gamma=0.001完全不符合常规取值，gamma作为折扣因子，通常应设置在0.9~0.999区间。过小的gamma会导致时序差分计算出现数值不稳定，这大概率是NaN出现的核心原因之一，先改回合理值：
model = MaskablePPO(MaskableActorCriticPolicy, env, gamma=0.99, verbose=0)
2. 归一化观测空间
自定义环境的观测值如果数值范围过大，会引发梯度爆炸进而产生NaN。强制将所有观测特征缩放到[0,1]或[-1,1]区间：
# 在自定义env的reset/step方法中处理观测值 def normalize_obs(self, obs): # 提前统计观测值的最大最小值，替换为你的实际数值 obs_min = np.array([0, 0, -10]) obs_max = np.array([100, 50, 10]) # 加1e-8避免除零错误 return (obs - obs_min) / (obs_max - obs_min + 1e-8)
3. 限制奖励范围
如果奖励值波动极大（比如突然出现1e4或-1e4的奖励），会导致值函数计算异常。对奖励做裁剪或缩放：
# 在自定义env的step方法中修改奖励返回逻辑 reward = np.clip(reward, -10, 10) # 限制奖励在[-10,10]区间 # 或做归一化 reward = reward / 100.0
4. 添加梯度裁剪
通过policy_kwargs配置梯度裁剪，限制权重更新的幅度，避免梯度爆炸：
model = MaskablePPO( MaskableActorCriticPolicy, env, gamma=0.99, verbose=0, policy_kwargs={ "optimizer_kwargs": {"clipnorm": 1.0} # 将梯度范数限制为1.0 } )
5. 校验动作掩码有效性
确保valid_action_mask()返回的掩码不会出现全False的情况（无有效动作可选），这种情况会导致logits计算异常：
def valid_action_mask(self): mask = ... # 你的掩码生成逻辑 # 兜底校验：如果没有有效动作，强制设置第一个动作有效 if not np.any(mask): mask[0] = True return mask
6. 使用数值稳定的激活函数
将神经网络的激活函数从ReLU改为LeakyReLU，避免神经元死亡导致的梯度消失/NaN：
import torch.nn as nn model = MaskablePPO( MaskableActorCriticPolicy, env, gamma=0.99, verbose=0, policy_kwargs={ "net_arch": dict(pi=[64, 64], vf=[64, 64]), "activation_fn": nn.LeakyReLU } )
内容的提问来源于stack exchange，提问作者Vahid Gholamzade
相关产品推荐
相关产品推荐
