ray.rllib.core.learner.learner_group.LearnerGroup.load_module_state#

LearnerGroup.load_module_state(*, multi_rl_module_ckpt_dir: str | None = None, modules_to_load: Set[str] | None = None, rl_module_ckpt_dirs: Dict[str, str] | None = None) → None[source]#

Load the checkpoints of the modules being trained by LearnerGroup.

load_module_state can be used 3 ways:
  1. Load a checkpoint for the MultiRLModule being trained by this

    LearnerGroup. Optionally, limit the modules that are loaded from the checkpoint by specifying the modules_to_load argument.

  2. Load the checkpoint(s) for single agent RLModules that

    are in the MultiRLModule being trained by this LearnerGroup.

  3. Load a checkpoint for the MultiRLModule being trained by this

    LearnerGroup and load the checkpoint(s) for single agent RLModules that are in the MultiRLModule. The checkpoints for the single agent RLModules take precedence over the module states in the MultiRLModule checkpoint.

At least one of multi_rl_module_ckpt_dir or rl_module_ckpt_dirs must be specified. modules_to_load can only be specified if multi_rl_module_ckpt_dir is provided.

Parameters:

  • multi_rl_module_ckpt_dir – The path to the checkpoint for the MultiRLModule.

  • modules_to_load – A set of RLModule ids to load from the checkpoint.

  • rl_module_ckpt_dirs – A mapping from module ids to the path to a checkpoint for a single agent RLModule.