Skip to content

RllibScholaModel

Full path: schola.rllib.export.RllibScholaModel

Wrapper for RLlib policies to enable ONNX export for use with Schola.

RllibScholaModel(policy, observation_space = None, action_space = None)

Bases: ScholaModel

Parameters

  • policy (ray.rllib.policy.Policy) - The RLlib policy to wrap for ONNX export.

  • observation_space (Space)

  • action_space (Space)

Methods

init

__init__(policy, observation_space = None, action_space = None)

Parameters

  • policy (Any)

  • observation_space (Space)

  • action_space (Space)

forward

forward(*args)

Forward pass through the model. Removes variance outputs, to make compatible with Unreal.

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

_policy

_policy : ray.rllib.policy.Policy

The wrapped RLlib policy.

_model

_model : torch.nn.Module

The underlying PyTorch model from the policy.

initial_state_dict

initial_state_dict

Nested structure of state tensors without batch dimensions.

state_metadata

state_metadata: Metadata aligned with initial_state_dict keys (default

empty StateMetadata for each leaf).