OakEmu-RL v0.1.0

Appearance

PPOModel ​

python
PPOModel(observation_shape: tuple[int, ...], action_shape: int, model_config: PPOModelConfiguration | None = None) -> None

Actor-critic model for PPO.

Uses a shared convolutional feature extractor followed by separate policy and value heads. The policy head outputs action logits; the value head outputs a scalar state value.

Initialize the PPOModel.

Parameters:

  • observation_shape (tuple[int, ...]) – The shape of the environment observations.
  • action_shape (int) – The number of possible actions.
  • model_config (PPOModelConfiguration | None) – Optional configuration for the model architecture.

Source

Methods ​

forward ​

python
forward(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]

Perform a forward pass through the model.

Parameters:

  • x (Tensor) – The input tensor of shape (batch, *observation_shape) in the source observation dtype.

Returns:

  • tuple[Tensor, Tensor] – A tuple of (action_logits, values) where action_logits has shape (batch, actions) and values has shape (batch, 1).

Source

get_action_bundle ​

python
get_action_bundle(x: torch.Tensor, action: torch.Tensor | None = None, action_mask: torch.Tensor | None = None) -> PPOActionBundle

Get action, log probability, entropy, and value estimates for the given states.

If action is provided, computes the log probability and entropy for that action. Otherwise, samples a new action from the policy.

Parameters:

  • x (Tensor) – The input tensor of shape (batch, *observation_shape).
  • action (Tensor | None) – Optional pre-selected actions of shape (batch,).
  • action_mask (Tensor | None) – Optional boolean mask of shape (actions,) where True = allowed.

Returns:

  • PPOActionBundle – The sampled action bundle.

Source

get_logits ​

python
get_logits(x: torch.Tensor) -> torch.Tensor

Compute action logits only (no value computation).

Used for greedy action selection in Tester and for ONNX export.

Parameters:

  • x (Tensor) – The input tensor of shape (batch, *observation_shape).

Returns:

  • Tensor – Action logits of shape (batch, actions).

Source

get_value ​

python
get_value(x: torch.Tensor, *, detach_features: bool = False) -> torch.Tensor

Compute state values only (no policy computation).

Parameters:

  • x (Tensor) – The input tensor of shape (batch, *observation_shape).
  • detach_features (bool) – When True, the shared trunk features are detached so gradients reach only the value head. Used by the BC bootstrap's critic-calibration step, which must not reshape the trunk features BC is imprinting the policy on.

Returns:

  • Tensor – Values of shape (batch,).

Source

Values ​

action_shape ​

python
action_shape = action_shape

Source

conv ​

python
conv = nn.Sequential(*conv_layers)

Source

fc ​

python
fc = nn.Sequential(*fc_layers)

Source

policy_head ​

python
policy_head = nn.Linear(in_features, action_shape)

Source

value_head ​

python
value_head = nn.Linear(in_features, 1)

Source