Skip to content

ScholaRLModule

Full path: schola.rllib.export.ScholaRLModule

ScholaModel adapter around an RLlib RLModule for ONNX export.

ScholaRLModule(rl_module, observation_space = None, action_space = None,
               ignored_state_keys = ('critic', ))

Bases: ScholaModel

Parameters

  • rl_module (ray.rllib.core.rl_module.rl_module.RLModule) - Source module whose forward pass is wrapped.

  • observation_space (gymnasium.spaces.Space) - Override observation space; defaults to rl_module.observation_space.

  • action_space (gymnasium.spaces.Space) - Override action space; defaults to rl_module.action_space.

  • ignored_state_keys (tuple) - State dict keys to drop before export (default hides critic-only state).

Methods

init

__init__(rl_module, observation_space = None, action_space = None,
         ignored_state_keys = ('critic', ))

Parameters

  • rl_module (Any)

  • observation_space (Space)

  • action_space (Space)

  • ignored_state_keys (Tuple)

forward

forward(*args)

Parameters

  • args (Any)

get_logit_dimensions

get_logit_dimensions()

Get the flat dimensions of the action spaces. :returns: Flat size per action dict key (gymnasium.spaces.flatdim on each subspace). :rtype: Dict[str, int]

Returns

Return type: Dict[str, int]

make_box_output

make_box_output(logits, space_name)

Map logits to a gymnasium.spaces.Box action slice (identity for Box).

Parameters

  • logits (torch.Tensor) - Logits slice for space_name (typically shaped for one fundamental space).

  • space_name (str) - Key in action_space used only for symmetry with other make_* helpers.

Attributes

initial_state_dict

initial_state_dict

Nested structure of state tensors without batch dimensions.

num_state_inputs

num_state_inputs