[Refactor] Dynamic value-estimator registry across all loss modules

[vmoens]

Context

Today every loss in torchrl.objectives ships its own make_value_estimator body with a hand-coded if/elif chain over ValueEstimators.{TD0,TD1,GAE,TDLambda,VTrace}. The chain knows the class names, the default hyperparameters, and any per-estimator construction quirks (e.g. V-Trace needs the actor). Adding a new estimator therefore means touching ~15 loss files plus their tests.

This PR replaces that with a registry: estimators self-register via a decorator at class-definition time, and every loss dispatches through a single helper.

This is the registry half of #3748 — pulled out into its own PR per review feedback so that the MAPPO/IPPO/MultiAgentGAE/ValueNorm work in #3748 stays scoped to its features, and the registry refactor can be reviewed independently and extended to all losses (the original PR only converted PPO / A2C / Reinforce, leaving the headline benefit unfulfilled for the other ~12 loss classes).

What's new

torchrl.objectives.utils

• _VALUE_ESTIMATOR_REGISTRY: dict[ValueEstimators, _ValueEstimatorRegistryEntry]
• @register_value_estimator(value_type, *, default_kwargs=...) — class decorator that estimators apply at definition time.
• get_value_estimator_entry(value_type) — accepts both ValueEstimators enum members and lowercase string aliases (`"gae"`, `"vtrace"`, ...). Handy for hydra/yaml configs.
• build_value_estimator(loss, value_type, **hp) — low-level: looks up the class via the registry, merges entry.default_kwargs with the caller's hp, then calls cls.for_loss(loss, **merged).
• dispatch_value_estimator(loss, value_type, *, supported, tensor_keys, **hp) — high-level helper most losses use. It:
  1. validates value_type against the loss's supported set;
  2. seeds gamma from loss.gamma if present;
  3. builds the estimator via the registry;
  4. assigns loss._value_estimator / loss.value_type;
  5. calls set_keys(**tensor_keys).
• default_value_kwargs() is now a 1-line shim over the registry (back-compat preserved).

torchrl.objectives.value.advantages

• ValueEstimatorBase grows a for_loss(cls, loss_module, **hp) classmethod. Default behaviour: pick loss.critic_network, fall back to loss.value_network. If the caller (i.e. the loss) passes value_network=<module> in hp, that wins.
• VTrace overrides for_loss to also pluck the actor and handle the functional-actor deep-copy.
• TD0Estimator / TD1Estimator / TDLambdaEstimator / GAE / VTrace each carry an @register_value_estimator(...) decorator with their default kwargs.

Per-loss conversion

Every loss now declares a SUPPORTED_VALUE_ESTIMATORS tuple, and its make_value_estimator body collapses to a single dispatch_value_estimator(...) call:

Loss family	Supported set
PPOLoss, A2CLoss, ReinforceLoss	TD0, TD1, TDLambda, GAE, VTrace
DQNLoss, IQLLoss, DiscreteIQLLoss, DDPGLoss, TD3Loss, TD3BCLoss, SACLoss, DiscreteSACLoss, CQLLoss, DiscreteCQLLoss, REDQLoss, REDQLoss_deprecated, DreamerActorLoss, DreamerV3ActorLoss, CrossQLoss	TD0, TD1, TDLambda
QMixerLoss	TD0

The instance/class fast-path for power users who pass a ValueEstimatorBase directly stays at the top of every body.

Why a partial conversion would have been bad

The version in #3748 only converted PPO/A2C/Reinforce. That left the if/elif chains intact in 12 other losses, and the PR still had to hand-edit 8 test files to add new estimators to the "raises NotImplementedError" lists. The registry's whole pitch ("no edits in existing losses when adding a new estimator") only holds if every loss participates. This PR finishes the conversion so the next estimator addition — including MAGAE in #3748 — really does require zero edits to existing losses.

Verification

• New TestValueEstimatorRegistry suite (8 tests) in test/objectives/test_loss_module.py:
  • all built-ins are registered
  • string alias resolution (\"gae\" ↔ ValueEstimators.GAE)
  • error paths for unknown alias / non-enum type
  • third-party estimator registration via the decorator
  • dispatch_value_estimator rejects estimators outside the supported set
  • value_network=... override path
  • default_value_kwargs() matches the registry
• Existing test suites pass unchanged:
  • pytest test/objectives/test_ppo.py test/objectives/test_dqn.py test/objectives/test_ddpg.py test/objectives/test_sac.py test/objectives/test_iql.py test/objectives/test_cql.py test/objectives/test_dreamer.py — 4521 passed.

Stats

• ~700 lines deleted, ~600 added across 19 files.
• All registry indirection is in torchrl/objectives/utils.py (one file, ~120 lines).

Out of scope (for #3748 to add on top)

• MAGAE / MultiAgentGAE registration — lives in [Feature] MAPPOLoss + IPPOLoss + MultiAgentGAE + ValueNorm #3748 with the rest of the MAPPO feature.

Made with Cursor

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/3780

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Run pull request jobs on OSDC runners in shadow mode

✅ No Failures

As of commit 1e03f13 with merge base 258dfad ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.