ray.rllib.core.learner.learner.Learner.filter_param_dict_for_optimizer

Learner.filter_param_dict_for_optimizer(param_dict: Dict[Hashable, torch.Tensor | tf.Variable], optimizer: torch.optim.Optimizer | keras.optimizers.Optimizer) → Dict[Hashable, torch.Tensor | tf.Variable]

Reduces the given ParamDict to contain only parameters for given optimizer.

Parameters:

• param_dict – The ParamDict to reduce/filter down to the given optimizer. The returned dict will be a subset of param_dict only containing keys (param refs) that were registered together with optimizer (and thus that optimizer is responsible for applying gradients to).

• optimizer – The optimizer object to whose parameter refs the given param_dict should be reduced.

Returns:

A new ParamDict only containing param ref keys that belong to optimizer.