simulo.RewardSet

Helper to compose multiple reward components.

Sums weighted rewards from multiple components into a single tensor.

Example: reward_set = simulo.RewardSet([ simulo.reward.HeadingAlignment(weight=0.5), simulo.reward.UprightPosture(weight=0.1), ]) reward_set.initialize(robot, device, num_envs)

In get_rewards():

reward = reward_set.compute(prev_actions, terminated)

simulo.RewardSet(components: List[RewardComponent])

initialize()

RewardSet.initialize(robot: RobotProtocol, device: str, num_envs: int) -> None

Initialize all components.

Must be called in on_start() before using compute().

Args:

• robot — Robot instance
• device — Torch device string
• num_envs — Number of parallel environments

compute()

RewardSet.compute(
    prev_actions: torch.Tensor,
    terminated: Optional[torch.Tensor] = None,
) -> torch.Tensor

Compute sum of all weighted rewards.

Args:

• prev_actions — Previous actions tensor
• terminated — Termination flags or None

Returns:

Total reward tensor of shape (num_envs,)

compute_breakdown()

RewardSet.compute_breakdown(
    prev_actions: torch.Tensor,
    terminated: Optional[torch.Tensor] = None,
) -> Dict[str, torch.Tensor]

Compute rewards with per-component breakdown.

Useful for debugging and understanding reward contributions.

Args:

• prev_actions — Previous actions tensor
• terminated — Termination flags or None

Returns:

Dict mapping component names to their reward tensors