protomotions.agents.ppo.utils package#

PPO advantage, motion indexing, and exploration helpers.

Advantages#

Utility functions for PPO algorithm.

This module provides helper functions for PPO, including advantage computation using Generalized Advantage Estimation (GAE).

Key Functions:

  • discount_values: Compute GAE advantages from rewards and values

protomotions.agents.ppo.utils.advantages.discount_values(
mb_fdones,
mb_values,
mb_rewards,
mb_next_values,
gamma,
tau,
)[source]#

Compute Generalized Advantage Estimation (GAE) advantages.

Computes advantages using GAE-Lambda, which provides a bias-variance tradeoff for advantage estimation. Uses backwards iteration through the episode to compute bootstrapped advantages.

Parameters:

  • mb_fdones – Done flags (num_steps, num_envs). 1.0 = episode ended.

  • mb_values – Value predictions at each timestep (num_steps, num_envs).

  • mb_rewards – Rewards received at each timestep (num_steps, num_envs).

  • mb_next_values – Value predictions for next states (num_steps, num_envs).

  • gamma – Discount factor for future rewards (typically 0.99).

  • tau – GAE lambda parameter for bias-variance tradeoff (typically 0.95).

Returns:

Tensor of advantages with shape (num_steps, num_envs).

Example

>>> advantages = discount_values(dones, values, rewards, next_values, 0.99, 0.95)
>>> returns = advantages + values  # Can compute returns from advantages

Note

GAE-Lambda provides advantages that balance bias (low lambda) vs. variance (high lambda). Lambda=0 gives 1-step TD, lambda=1 gives Monte Carlo returns.

Reference:

Schulman et al. “High-Dimensional Continuous Control Using Generalized Advantage Estimation” (2015)

Motion indexing#

Training-only, stable row allocation for per-motion Gaussian exploration.

protomotions.agents.ppo.utils.motion_index.names_digest(names)[source]#
class protomotions.agents.ppo.utils.motion_index.MotionShard(
path: 'str',
offset: 'int',
count: 'int',
names_digest: 'str',
)[source]#
path: str#
offset: int#
count: int#
names_digest: str#
__init__(path, offset, count, names_digest)#
class protomotions.agents.ppo.utils.motion_index.MotionIndex(shards: 'list[MotionShard]')[source]#
shards: list[MotionShard]#
property num_motions#
state_dict()[source]#
offset_for(motion_lib)[source]#
classmethod from_sources(sources)[source]#

Read each packaged shard’s metadata once; mmap avoids loading poses.

classmethod build(motion_lib)[source]#
__init__(shards)#

Per-motion exploration#

Checkpoint preparation and inspection snapshots for per-motion exploration.

protomotions.agents.ppo.utils.per_motion_exploration.prepare_checkpoint(
checkpoint,
motion_index,
actor,
load_training_state,
)[source]#

Validate row identity/bounds or expand shared std for a weights-only start.

Neither the input checkpoint nor the actor is modified. Validation happens before the agent loads any weights or optimizer state.

protomotions.agents.ppo.utils.per_motion_exploration.save_snapshot(path, motion_index, logstd, epoch, step_count)[source]#

Atomically replace the inspection sidecar; resume uses the full checkpoint.

The agent owns rank-zero selection and distributed error propagation.