OptimizerHook

class torchrl.trainers.OptimizerHook(optimizer: Optimizer, loss_components: Optional[Sequence[str]] = None)

Add an optimizer for one or more loss components.

Parameters:

• optimizer (optim.Optimizer) – An optimizer to apply to the loss_components.

• loss_components (Sequence[str], optional) – The keys in the loss TensorDict for which the optimizer should be appled to the respective values. If omitted, the optimizer is applied to all components with the names starting with loss_.

Examples

>>> optimizer_hook = OptimizerHook(optimizer, ["loss_actor"])
>>> trainer.register_op("optimizer", optimizer_hook)

register(trainer, name='optimizer') → None

Registers the hook in the trainer at a default location.

Parameters:

• trainer (Trainer) – the trainer where the hook must be registered.

• name (str) – the name of the hook.

Note

To register the hook at another location than the default, use register_op().