ray.rllib.core.learner.learner.Learner.register_optimizer#

Learner.register_optimizer(*, module_id: str = '__all__', optimizer_name: str = 'default_optimizer', optimizer: Union[torch.optim.optimizer.Optimizer, keras.optimizers.optimizer_experimental.optimizer.Optimizer], params: Sequence[Union[torch.Tensor, tensorflow.python.ops.variables.Variable]], lr_or_lr_schedule: Optional[Union[float, List[List[Union[int, float]]]]] = None) → None[source]#

Registers an optimizer with a ModuleID, name, param list and lr-scheduler.

Use this method in your custom implementations of either self.configure_optimizers() or self.configure_optimzers_for_module() (you should only override one of these!). If you register a learning rate Scheduler setting together with an optimizer, RLlib will automatically keep this optimizer’s learning rate updated throughout the training process. Alternatively, you can construct your optimizers directly with a learning rate and manage learning rate scheduling or updating yourself.

Parameters

  • module_id – The module_id under which to register the optimizer. If not provided, will assume ALL_MODULES.

  • optimizer_name – The name (str) of the optimizer. If not provided, will assume DEFAULT_OPTIMIZER.

  • optimizer – The already instantiated optimizer object to register.

  • params – A list of parameters (framework-specific variables) that will be trained/updated

  • lr_or_lr_schedule – An optional fixed learning rate or learning rate schedule setup. If provided, RLlib will automatically keep the optimizer’s learning rate updated.