From 0f69759da48e76705c9247845d9dabff851e7b5f Mon Sep 17 00:00:00 2001 From: zb-zhoufengen Date: Wed, 20 May 2026 10:52:01 +0800 Subject: [PATCH] fix(ppo): protect against NaN/Inf gradient corruption in before_step NaN or Inf gradients can silently corrupt model weights when they reach optimizer.step() during PPO training. A single bad batch propagates NaN into all downstream weights, ruining multi-day training runs without any error signal. Add torch.nan_to_num_ after gradient clipping in before_step() to replace NaN with zero and clamp Inf to max_grad_norm bounds, preventing silent weight corruption. This is a backward-compatible safety net: valid gradients are unchanged, and the additional iteration over parameters is negligible compared to the backward pass. Fixes #2226 --- .../habitat_baselines/rl/ppo/ppo.py | 5 + test/test_ppo_nan_gradients.py | 193 ++++++++++++++++++ 2 files changed, 198 insertions(+) create mode 100644 test/test_ppo_nan_gradients.py diff --git a/habitat-baselines/habitat_baselines/rl/ppo/ppo.py b/habitat-baselines/habitat_baselines/rl/ppo/ppo.py index c99ec55f91..2bc0a47a1f 100644 --- a/habitat-baselines/habitat_baselines/rl/ppo/ppo.py +++ b/habitat-baselines/habitat_baselines/rl/ppo/ppo.py @@ -368,6 +368,11 @@ def before_step(self) -> torch.Tensor: [h.wait() for h in handles] + # Replace NaN/Inf gradients with zero to prevent silent weight corruption + for p in self.parameters(): + if p.grad is not None: + p.grad.data.nan_to_num_(nan=0.0, posinf=self.max_grad_norm, neginf=-self.max_grad_norm) + return grad_norm def after_step(self) -> None: diff --git a/test/test_ppo_nan_gradients.py b/test/test_ppo_nan_gradients.py new file mode 100644 index 0000000000..b6189c4ca5 --- /dev/null +++ b/test/test_ppo_nan_gradients.py @@ -0,0 +1,193 @@ +#!/usr/bin/env python3 + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +""" +Tests for NaN/Inf gradient protection in PPO.before_step(). + +A single bad batch can produce NaN gradients that silently corrupt +model weights through optimizer.step(). The before_step() hook in +PPO now replaces NaN/Inf with safe values before the optimizer step. + +This test duplicates the before_step() logic in a minimal PPO-like +test class so it runs without habitat-sim (conda-only dependency). +""" + +import pytest +import torch +import torch.nn as nn +import torch.optim as optim + + +class MockActorCritic(nn.Module): + """Minimal actor-critic stub matching NetPolicy interface.""" + + def __init__(self): + super().__init__() + self.policy = nn.Linear(8, 4) + + def policy_parameters(self): + return self.policy.parameters() + + def aux_loss_parameters(self): + return {} + + +class PPOBeforeStep: + """Test-only surrogate that replicates PPO.before_step() logic. + + The real before_step() in habitat_baselines does: + 1. Distributed all-reduce for non-actor-critic params + 2. clip_grad_norm_ on policy parameters + 3. clip_grad_norm_ on aux loss parameters + 4. Wait for all-reduce handles + 5. Replace NaN/Inf gradients with safe values <-- the fix under test + """ + + def __init__(self, max_grad_norm=0.5): + self.actor_critic = MockActorCritic() + self.max_grad_norm = max_grad_norm + # non_ac_params: params whose name does not start with "actor_critic." + self.non_ac_params = [] + + # Create optimizer so parameter gradients are initialized + self.optimizer = optim.SGD( + self.parameters(), lr=0.001 + ) + + # Initialize all gradients to zero + self._zero_grads() + + def _zero_grads(self): + for p in self.parameters(): + p.grad = torch.zeros_like(p.data) + + def parameters(self): + return self.actor_critic.parameters() + + def named_parameters(self): + return self.actor_critic.named_parameters() + + def before_step(self): + """Replicates the exact logic from PPO.before_step() in ppo.py.""" + handles = [] + if torch.distributed.is_initialized(): + for p in self.non_ac_params: + if p.grad is not None: + p.grad.data.detach().div_( + torch.distributed.get_world_size() + ) + handles.append( + torch.distributed.all_reduce( + p.grad.data.detach(), async_op=True + ) + ) + + grad_norm = nn.utils.clip_grad_norm_( + self.actor_critic.policy_parameters(), + self.max_grad_norm, + ) + + for v in self.actor_critic.aux_loss_parameters().values(): + nn.utils.clip_grad_norm_(v, self.max_grad_norm) + + [h.wait() for h in handles] + + # Replace NaN/Inf gradients with zero to prevent silent weight corruption + for p in self.parameters(): + if p.grad is not None: + p.grad.data.nan_to_num_( + nan=0.0, + posinf=self.max_grad_norm, + neginf=-self.max_grad_norm, + ) + + return grad_norm + + +@pytest.fixture +def updater(): + """Create a PPOBeforeStep instance for testing.""" + return PPOBeforeStep(max_grad_norm=0.5) + + +def test_nan_gradients_replaced_with_zero(updater): + """NaN gradients are replaced with zero in before_step().""" + param = next(updater.parameters()) + param.grad = torch.full(param.grad.shape, float("nan")) + + updater.before_step() + + assert not torch.isnan(param.grad).any(), "NaN gradients were not replaced" + + +def test_inf_gradients_clamped_to_max_norm(updater): + """Inf gradients are clamped to max_grad_norm bounds.""" + param = next(updater.parameters()) + param.grad = torch.full(param.grad.shape, float("inf")) + + updater.before_step() + + assert not torch.isinf(param.grad).any(), "Positive inf was not clamped" + + param.grad = torch.full(param.grad.shape, float("-inf")) + updater.before_step() + + assert not torch.isinf(param.grad).any(), "Negative inf was not clamped" + + +def test_valid_gradients_unchanged(updater): + """Finite gradients pass through before_step() unmodified.""" + param = next(updater.parameters()) + original = param.grad.clone() + orig_values = original.flatten()[:6].clone() + # Set some non-zero finite values + param.grad.flatten()[:6] = torch.tensor([0.1, -0.2, 0.05, -0.01, 0.3, -0.15]) + original = param.grad.clone() + + updater.before_step() + + assert torch.allclose(param.grad, original), ( + "Valid gradients were incorrectly modified" + ) + + +def test_all_parameters_cleaned(updater): + """NaN/Inf replacement applies to every parameter.""" + for p in updater.parameters(): + if p.grad is not None: + p.grad.fill_(float("nan")) + + updater.before_step() + + for p in updater.parameters(): + if p.grad is not None: + assert not torch.isnan(p.grad).any(), ( + "NaN remains in gradient after before_step" + ) + assert not torch.isinf(p.grad).any(), ( + "Inf remains in gradient after before_step" + ) + + +def test_optimizer_step_with_nan_gradients(updater): + """End-to-end: NaN gradients do not corrupt model weights after step(). + + This is the core scenario from issue #2226: a single batch with + NaN gradients should not permanently corrupt the optimizer state. + """ + # Inject NaN gradients (simulates a bad loss.backward() call) + for p in updater.parameters(): + if p.grad is not None: + p.grad.fill_(float("nan")) + + updater.before_step() + updater.optimizer.step() + + # Verify weights are still finite + for name, p in updater.named_parameters(): + assert torch.isfinite(p).all(), ( + f"Weights for {name} became NaN/Inf after optimizer.step()" + )