ray.rllib.core.rl_module.rl_module.RLModule

class ray.rllib.core.rl_module.rl_module.RLModule(config: ray.rllib.core.rl_module.rl_module.RLModuleConfig)

Bases: abc.ABC

Base class for RLlib modules.

Subclasses should call super().__init__(config) in their __init__ method. Here is the pseudocode for how the forward methods are called:

During Training (acting in env from each rollout worker):

module = RLModule(...)
obs, info = env.reset()

while not env.terminated:
    fwd_outputs = module.forward_exploration({"obs": obs})
    # this can be either deterministic or stochastic distribution
    action = fwd_outputs["action_dist"].sample()
    obs, reward, terminated, truncated, info = env.step(action)

During Training (learning the policy)

module = RLModule(...)
fwd_ins = {"obs": obs, "action": action, "reward": reward, "next_obs": next_obs}
fwd_outputs = module.forward_train(fwd_ins)
loss = compute_loss(fwd_outputs, fwd_ins)
update_params(module, loss)

During Inference (acting in env during evaluation)

module = RLModule(...)
obs, info = env.reset()

while not env.terminated:
    fwd_outputs = module.forward_inference({"obs": obs})
    action = fwd_outputs["action_dist"].sample()
    obs, reward, terminated, truncated, info = env.step(action)

Parameters

config – The config for the RLModule.

Abstract Methods:

forward_train(): Forward pass during training.

forward_exploration(): Forward pass during training for exploration.

forward_inference(): Forward pass during inference.

Note

There is a reason that the specs are not written as abstract properties. The reason is that torch overrides __getattr__ and __setattr__. This means that if we define the specs as properties, then any error in the property will be interpreted as a failure to retrieve the attribute and will invoke __getattr__ which will give a confusing error about the attribute not found. More details here: https://github.com/pytorch/pytorch/issues/49726.

Methods

as_multi_agent()	Returns a multi-agent wrapper around this module.
forward_exploration(batch, **kwargs)	Forward-pass during exploration, called from the sampler.
forward_inference(batch, **kwargs)	Forward-pass during evaluation, called from the sampler.
forward_train(batch, **kwargs)	Forward-pass during training called from the learner.
from_checkpoint(checkpoint_dir_path)	Loads the module from a checkpoint directory.
get_exploration_action_dist_cls()	Returns the action distribution class for this RLModule used for exploration.
get_inference_action_dist_cls()	Returns the action distribution class for this RLModule used for inference.
get_initial_state()	Returns the initial state of the module.
get_state()	Returns the state dict of the module.
get_train_action_dist_cls()	Returns the action distribution class for this RLModule used for training.
input_specs_exploration()	Returns the input specs of the forward_exploration method.
input_specs_inference()	Returns the input specs of the forward_inference method.
input_specs_train()	Returns the input specs of the forward_train method.
load_state(dir)	Loads the weights of an RLModule from the directory dir.
output_specs_exploration()	Returns the output specs of the forward_exploration method.
output_specs_inference()	Returns the output specs of the forward_inference method.
output_specs_train()	Returns the output specs of the forward_train method.
save_state(dir)	Saves the weights of this RLModule to the directory dir.
save_to_checkpoint(checkpoint_dir_path)	Saves the module to a checkpoint directory.
set_state(state_dict)	Sets the state dict of the module.
setup()	Sets up the components of the module.
unwrapped()	Returns the underlying module if this module is a wrapper.

Attributes

framework