Bug Fix: sanitise non-finite gradients in PPO.before_step (#2226)

[dparikh79]

Summary

Fixes #2226.

torch.nn.utils.clip_grad_norm_ is NaN-unsafe: a single NaN in any parameter's gradient produces a NaN total_norm, a NaN clip_coef, and silently leaves NaN gradients in place. optimizer.step() then writes NaN into every parameter. The training loop keeps running with a dead model until something like torch.multinomial finally raises on a NaN probability tensor, by which point recent checkpoints contain NaN parameters.

In DD-PPO this is especially bad: DDP all-reduce averages NaN with finite into NaN, so a single bad mini-batch on one worker corrupts every worker's model on the next step. @alunxu's issue write-up walks through a reproduction on a ~180M-frame DD-PPO PointGoal navigation run where this happened in production.

Change

habitat-baselines/habitat_baselines/rl/ppo/ppo.py, PPO.before_step: before clip_grad_norm_ runs, sanitise non-finite elements in every parameter gradient with torch.nan_to_num_(p.grad, nan=0.0, posinf=0.0, neginf=0.0). This turns a NaN mini-batch into a no-op update instead of a permanent weight corruption.

The fix runs nan_to_num_ only when torch.isfinite(p.grad).all() is False, so:

• Clean training paths: nan_to_num_ of a finite tensor is identity, and we skip the call entirely. Bitwise-identical weights, no measurable perf hit.
• NaN-event paths: zero-clamped gradient, clip_grad_norm_ returns a finite number, optimizer.step() updates with lr * 0 = 0, model weights unchanged.

The fix matches @alunxu's exact proposal in the issue.

Minimal reproduction (no habitat needed)

Borrowed from @alunxu's #2226 write-up; demonstrates that today's clip_grad_norm_ silently leaves NaN gradients and optimizer.step() corrupts weights:

import torch
import torch.nn as nn

model = nn.Linear(4, 2)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)

# Simulate a NaN gradient (e.g. from log(0) in policy entropy)
for p in model.parameters():
    p.grad = torch.full_like(p, float("nan"))

# This is what PPO.before_step does today:
norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.2)
print("grad norm:", norm)                              # nan
print("grads finite:", all(torch.isfinite(p.grad).all() for p in model.parameters()))   # False
opt.step()
print("weights finite:", all(torch.isfinite(p).all() for p in model.parameters()))      # False

Applying the sanitise step before clip_grad_norm_ flips both grads finite and weights finite to True.

Test plan

• Read ppo.py:347-371 to confirm the patched function is the one DD-PPO actually calls via the before_step hook.
• Confirmed via grep that clip_grad_norm_ is the only path that mutates p.grad before optimizer.step() in this code path.
• Could not run habitat-lab tests locally: the existing PPO/DDPPO tests (test_ddppo_reduce.py, test_pointnav_resnet_policy.py) require the full habitat env + scene data downloads + distributed setup. Adding an isolated unit test for the NaN-sanitise behaviour without that infrastructure would either require extracting the sanitise block into a helper (larger surface) or building a minimal PPO instance (heavy). Happy to add one in a follow-up if a maintainer prefers a particular shape.

CLA

Will need to sign the Meta CLA on first PR; will do so when the bot prompts.

Acknowledgements

• @alunxu for the rigorous bug analysis, root cause walk, and minimal reproduction in Silent weight corruption: PPO.before_step lets NaN gradients reach optimizer.step() #2226.
• @wadeKeith for the prior interest noted on the issue (their personal fork PR link in the comment is no longer reachable).

AI Assistance Disclosure

Implementation drafted with Claude assistance against @alunxu's proposed fix. I reviewed every changed line, traced the call path through before_step, and confirmed that the conditional if not torch.isfinite(p.grad).all() guard makes the change a no-op on clean training paths. I am the human contributor accountable for this PR.

[meta-cla]

Hi @dparikh79!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[dparikh79]

@facebook-github-bot recheck cla

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!