Skip to content

RewardShapingPipeline ​

python
RewardShapingPipeline(reward_shapers: RewardShaperConfiguration, context: RewardShapingContext) -> None

Sequential trainer-owned reward-shaping pipeline.

Initialize the reward-shaping pipeline.

Parameters:

  • reward_shapers (RewardShaperConfiguration) – None, a single shaper type, or a list of shaper types.
  • context (RewardShapingContext) – Shared shaping runtime context.

Source

Properties ​

has_intrinsic_reward_stream ​

python
has_intrinsic_reward_stream: bool

Return whether any configured shaper contributes to the intrinsic reward stream.

Returns:

  • bool – True when at least one shaper targets the intrinsic stream.

Source

is_enabled ​

python
is_enabled: bool

Return whether any reward shapers are configured.

Returns:

  • bool – True when at least one shaper is active.

Source

Methods ​

complete_episode ​

python
complete_episode(env_id: int, step: int) -> None

Log and clear reward-shaping episode metrics for one environment.

Parameters:

  • env_id (int) – The environment ID.
  • step (int) – The global step at episode end.

Source

flush_metrics ​

python
flush_metrics(step: int) -> None

Emit any throttled actor or learner reward-shaping metrics.

Parameters:

  • step (int) – The step at which to log.

Source

learn ​

python
learn(batch: RewardShapingBatch) -> None

Run learner-side reward-shaper updates for one batch.

Parameters:

  • batch (RewardShapingBatch) – The learner batch to consume.

Source

load_state_dict ​

python
load_state_dict(state: dict[str, Any]) -> None

Restore pipeline state from a checkpoint.

Parameters:

  • state (dict[str, Any]) – The serialised pipeline state.

Source

on_episode_start ​

python
on_episode_start(env_id: int, initial_observation: np.ndarray) -> None

Notify all configured shapers about a reset.

Parameters:

  • env_id (int) – The environment ID.
  • initial_observation (ndarray) – The reset observation.

Source

restore_from_checkpoint ​

python
restore_from_checkpoint(checkpoint: Checkpoint) -> None

Restore reward-shaper state from a checkpoint.

Parameters:

  • checkpoint (Checkpoint) – The checkpoint to restore from.

Raises:

  • OSError – If the checkpoint contents do not match the configured reward shapers.

Source

shape ​

python
shape(step: RewardShapingStep) -> RewardShapingResult

Apply all configured shapers in sequence and log reward metrics.

Parameters:

  • step (RewardShapingStep) – The step information to shape.

Returns:

  • RewardShapingResult – The shaped reward result.

Source

shape_batch ​

python
shape_batch(steps: Sequence[RewardShapingStep]) -> list[RewardShapingResult]

Apply all configured shapers to a batch of steps.

Shapers are run sequentially, but each individual shaper can override RewardShaper.shape_batch to do efficient batched inference.

Parameters:

  • steps (Sequence[RewardShapingStep]) – The steps to shape.

Returns:

  • list[RewardShapingResult] – The shaped reward results in the same order.

Source

state_dict ​

python
state_dict() -> dict[str, Any]

Capture pipeline state for checkpointing.

Returns:

  • dict[str, Any] – A serializable state dictionary.

Source

to_checkpoint_file ​

python
to_checkpoint_file(compress: bool = True) -> CheckpointFile | None

Build the checkpoint file for this pipeline.

Parameters:

  • compress (bool) – Whether to zstd-compress the checkpoint file.

Returns:

  • CheckpointFile | None – The checkpoint file, or None when reward shaping is disabled.

Source