Improve ChaosGrad diagnostics correctness and decay stability

[Copilot]

Summary

This PR tightens ChaosGrad behavior in two places: diagnostics now reflect effective per-group genesis LR overrides, and decoupled weight decay is numerically guarded under extreme LR×decay combinations.

Changes

• Diagnostics correctness
  • ChaosGrad.get_diagnostics() now computes avg_init_lr from each param group’s current lr (not only optimizer defaults).
  • Added avg_init_lr to each debug param_groups entry for group-level visibility.
• Numerical stability
  • In step(), decay multiplier is now floored with self._EPS:

    decay_factor = max(self._EPS, 1.0 - genesis_lr * per_decay)
    p.data.mul_(decay_factor)

  • Prevents pathological zero/negative scaling when genesis_lr * per_decay is large.
• Test coverage
  • Extended ChaosGrad extra tests to verify LR override propagation in diagnostics at both:
    • global level (avg_init_lr)
    • per-group level (param_groups[*].avg_init_lr)

Checklist

All contributions

• Code follows the project conventions (see CONTRIBUTING.md)
• Tests pass (python -m pytest tests/or pytest tests/)
• No sys.path.append hacks — imports use from odyssnet import ... directly

Library changes (odyssnet/)

• Corresponding test added/updated under tests/
• Documentation updated in relevant markdown files (LIBRARY.md, CONTRIBUTING.md)

New/modified example scripts (examples/)

• set_seed(42) is called as the first line of main()
• Placed in correct folder (examples/ for core validations, examples/advanced/ for complex tasks)
• Uses OdyssNetTrainer (not a manual training loop)
• Uses TrainingHistory to record metrics and calls history.plot() at the end
• File paths are relative to __file__, not hardcoded

[Copilot]

Pull request overview

This PR updates the ChaosGrad optimizer to (1) make diagnostics reflect per-parameter-group LR overrides rather than only optimizer defaults, and (2) add a numerical safety floor to the decoupled weight decay multiplier under extreme LR×decay values.

Changes:

• Compute avg_init_lr in get_diagnostics() using each param group’s effective lr, and include avg_init_lr in per-group debug diagnostics.
• Floor the decoupled weight decay multiplier in step() to avoid negative/near-zero scaling under extreme settings.
• Add tests ensuring diagnostics reflect LR overrides at both global and per-group levels.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
odyssnet/training/chaos_optimizer.py	Adjusts decay multiplier guarding and corrects diagnostics to use per-group genesis LR overrides.
tests/training/test_chaos_optimizer_extra.py	Adds coverage ensuring LR overrides propagate into global and per-group diagnostics.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

decay_factor is floored with self._EPS (1e-8), but when parameters are float16 this value can underflow to 0 when cast during mul_, reintroducing the pathological “zero scaling” this guard is trying to prevent. Consider using a dtype-aware floor (e.g., max(self._EPS, torch.finfo(p.data.dtype).tiny) or computing the factor as a tensor in the param dtype) so the clamp is effective under AMP/mixed precision as well.

Suggested change

	decay_factor = max(self._EPS, 1.0 - genesis_lr * per_decay)
	decay_floor = max(self._EPS, torch.finfo(p.data.dtype).tiny)
	decay_factor = max(decay_floor, 1.0 - genesis_lr * per_decay)

[Copilot]

genesis_lr = self.defaults['lr'] is now unused after switching avg_init_lr to be computed from each group's lr. Please remove the dead assignment to avoid confusion about which LR is actually being diagnosed.

Suggested change

genesis_lr = self.defaults['lr']