Introduce a Custom Algorithm
SUMMARY
The runner owns the loop — collect, bootstrap, update, log, checkpoint — and calls into the algorithm through a fixed set of methods. Anything implementing those methods can take PPO's place.
In practice, subclassing PPO and overriding one method is the shortest path, because the rollout storage, the normalizers and the checkpoint format come along with it.
The protocol the runner drives
Per iteration, OnPolicyRunner.learn() calls exactly this:
# collection, num_steps_per_env times
actions = alg.act(obs)
obs, rewards, dones, extras = env.step(actions)
alg.process_env_step(obs, rewards, dones, extras)
# then, once
alg.compute_returns(obs)
loss_dict = alg.update()and reads alg.learning_rate, alg.get_action_std() and alg.diagnostics for logging. Around that it calls alg.train_mode(), alg.eval_mode(), alg.compile(mode), alg.save(), alg.load(...) and alg.get_policy().
| Member | Contract |
|---|---|
act(obs) | Returns actions (num_envs, num_actions). Also where the transition's value and log-probability are recorded |
process_env_step(obs, rewards, dones, extras) | Records one transition, updates the observation normalizers, bootstraps on extras["time_outs"], and adds the transition to storage |
compute_returns(obs) | Computes returns and advantages from the final observation of the rollout |
update() | Optimizes, clears storage, and returns a dict of scalar losses. Keys appear as Loss/<key> |
train_mode() / eval_mode() | Switch the models |
save() / load(dict, cfg, strict) | The algorithm's slice of a checkpoint |
get_policy() | The actor, which is what gets exported |
get_action_std() | Reported as Policy/action_std, or None |
learning_rate | Attribute, reported as Learning/learning_rate |
diagnostics | Dict of scalars from the last update. Keys appear as Diagnostics/<key> |
compile(mode) | Applies torch.compile at the requested mode, or does nothing |
broadcast_parameters() / reduce_parameters() | Multi-GPU only |
The runner constructs it with keyword arguments:
alg = alg_class(
cfg=alg_cfg, # your algorithm config
actor=actor, # already sized for the observation set it consumes
critic=critic,
storage=storage, # a RolloutBuffer, already sized
obs_groups=obs_groups,
device=device,
multi_gpu_cfg=multi_gpu_cfg,
env=env,
)Note what the runner has already done for you: it built and sized the actor and critic from the observation groups, allocated the rollout buffer, and resolved the device. Your algorithm receives them rather than creating them.
The short path: subclass PPO
Most changes are one method. Everything else — storage, normalizers, checkpoint format, multi-GPU, compile — comes for free.
import torch
from telekinesis.rlbotics.algorithms import PPO
class EntropyDecayPPO(PPO):
"""PPO with an entropy coefficient that decays as the policy improves."""
def update(self) -> dict[str, float]:
# Anneal before the optimization the base class performs
self.cfg = self.cfg.replace(
entropy_coef=max(self.cfg.entropy_coef * 0.999, 1e-4)
)
loss_dict = super().update()
# Anything added here shows up as Loss/<key> in TensorBoard
loss_dict["entropy_coef"] = self.cfg.entropy_coef
return loss_dictWire it in by passing the class itself as class_name:
from telekinesis.rlbotics.config import OnPolicyRunnerConfig, PPOConfig
cfg = OnPolicyRunnerConfig(
obs_groups={"actor": ["policy"], "critic": ["policy"]},
algorithm=PPOConfig(class_name=EntropyDecayPPO, learning_rate=1e-3),
)Pass the class, not a string
A string class_name is resolved through the ALGORITHMS registry, which contains only {"PPO": PPO} — so "my_pkg.algorithms:MyPPO" raises a KeyError at construction. A callable is used as-is, which is why the class object is the working form.
One consequence: config.json records a class object through str(), as <class 'my_pkg.MyPPO'>, which from_dict cannot read back. A run configured with a custom algorithm is reproducible from your script, not from its config dump.
Adding hyperparameters
PPOConfig.from_dict rejects unknown keys, so a new hyperparameter needs its own config. Subclass PPOConfig — the runner accepts any subclass — and extend expected_class_names so validation admits your class name:
from dataclasses import dataclass
from typing import ClassVar
from telekinesis.rlbotics.config import PPOConfig
@dataclass
class EntropyDecayPPOConfig(PPOConfig):
"""PPOConfig plus the decay rate."""
expected_class_names: ClassVar[tuple[str, ...]] = ("PPO", "EntropyDecayPPO")
entropy_decay: float = 0.999
entropy_floor: float = 1e-4
def __post_init__(self) -> None:
super().__post_init__()
if not 0.0 < self.entropy_decay <= 1.0:
raise ValueError(f"'entropy_decay' must be in (0, 1], got {self.entropy_decay}.")Your algorithm then reads self.cfg.entropy_decay, and the config still validates itself and still converts to and from dictionaries. Follow the library's own convention and validate in __post_init__, so a bad value fails where it was written rather than mid-run.
WARNING
The runner requires the algorithm config to be a PPOConfig or a subclass of it. A config that does not inherit from it is rejected with Unsupported algorithm config type.
Writing one from scratch: SAC
Off-policy is the interesting case, because it is where the on-policy shape stops fitting for free. SAC is worth walking through because every seam shows up.
What the runner gives you, and what it does not
| The runner hands you | Usable for SAC? |
|---|---|
actor — a stochastic policy, num_actions out | Yes. Use SquashedGaussianDistribution so actions are bounded, which SAC assumes |
critic — a deterministic scalar, obs in | No. SAC needs twin Q(s, a), so the input is obs plus action. Build your own and leave this one unused |
storage — a RolloutBuffer sized num_steps_per_env | No. It is cleared at the end of every update(), so it cannot hold a replay history. Bring your own |
obs_groups, device, env, multi_gpu_cfg | Yes, unchanged |
The collection loop still calls act then process_env_step for num_steps_per_env steps, then compute_returns, then update once. SAC maps onto that as collect a chunk, then do N gradient steps from the replay buffer — which is how most parallel-environment SAC implementations behave anyway. compute_returns becomes a no-op: GAE has no meaning off-policy.
Pre-allocate your replay buffer, do not append to a list
Collection runs inside torch.inference_mode(), so every tensor produced there is an inference tensor and cannot later be saved for backward. update() runs outside it, so a buffer of appended inference tensors blows up with Inference tensors cannot be saved for backward the first time a Q network differentiates one.
RolloutBuffer avoids this by allocating its tensors up front and copy_-ing into them, and your replay buffer has to do the same.
The config
The runner rejects an algorithm config that is not a PPOConfig or a subclass, so SAC's config subclasses it. You inherit fields SAC does not use — clip_param, lam, desired_kl — which is the price of that check.
from dataclasses import dataclass
from telekinesis.rlbotics.config import PPOConfig
@dataclass
class SACConfig(PPOConfig):
"""PPOConfig plus what SAC needs. Inherited PPO-only fields are ignored."""
tau: float = 0.005 # Polyak coefficient for the target critics
init_alpha: float = 0.2 # Initial entropy temperature
learn_alpha: bool = True # Tune the temperature against a target entropy
target_entropy: float | None = None # Defaults to -num_actions when None
replay_capacity: int = 1_000_000 # Transitions kept, across all environments
batch_size: int = 256
gradient_steps: int = 64 # Updates per learning iteration
learning_starts: int = 1_000 # Transitions to collect before updating at all
def __post_init__(self) -> None:
super().__post_init__()
if not 0.0 < self.tau <= 1.0:
raise ValueError(f"'tau' must be in (0, 1], got {self.tau}.")
if self.init_alpha <= 0.0:
raise ValueError(f"'init_alpha' must be positive, got {self.init_alpha}.")
for name in ("replay_capacity", "batch_size", "gradient_steps"):
if getattr(self, name) <= 0:
raise ValueError(f"'{name}' must be positive, got {getattr(self, name)}.")The algorithm
import copy
import torch
from tensordict import TensorDict
from telekinesis.rlbotics.config import MLPConfig
from telekinesis.rlbotics.models import MLPModel
class ReplayBuffer:
"""A flat, pre-allocated replay buffer over all environments."""
def __init__(self, capacity, num_envs, obs_dim, num_actions, device):
self.capacity, self.num_envs, self.device = capacity, num_envs, device
shape = (capacity, obs_dim)
# Allocated here, outside inference mode, so update() can differentiate what it reads
self.obs = torch.zeros(shape, device=device)
self.next_obs = torch.zeros(shape, device=device)
self.actions = torch.zeros(capacity, num_actions, device=device)
self.rewards = torch.zeros(capacity, 1, device=device)
self.dones = torch.zeros(capacity, 1, device=device)
self.position, self.size = 0, 0
def add(self, obs, actions, rewards, next_obs, dones):
"""Copy one step from every environment in, wrapping around at capacity."""
index = (torch.arange(self.num_envs, device=self.device) + self.position) % self.capacity
self.obs[index] = obs
self.actions[index] = actions
self.rewards[index] = rewards.view(-1, 1)
self.next_obs[index] = next_obs
self.dones[index] = dones.view(-1, 1)
self.position = int((self.position + self.num_envs) % self.capacity)
self.size = min(self.size + self.num_envs, self.capacity)
def sample(self, batch_size):
index = torch.randint(0, self.size, (batch_size,), device=self.device)
return self.obs[index], self.actions[index], self.rewards[index], self.next_obs[index], self.dones[index]
class SAC:
"""Soft Actor-Critic against the runner's algorithm protocol."""
def __init__(self, cfg, actor, critic, storage, obs_groups=None, device="cpu",
multi_gpu_cfg=None, env=None):
self.cfg, self.device, self.obs_groups = cfg, device, obs_groups
self.actor = actor.to(device)
# The runner's critic is V(s); SAC needs Q(s, a), so build twin critics of our own
obs_dim = actor.input_dim
num_actions = env.num_actions
q_cfg = MLPConfig(hidden_dims=(256, 256), activation="elu", obs_normalization=False)
self.q1 = MLPModel(q_cfg, input_dim=obs_dim + num_actions, output_dim=1).to(device)
self.q2 = MLPModel(q_cfg, input_dim=obs_dim + num_actions, output_dim=1).to(device)
self.q1_target = copy.deepcopy(self.q1)
self.q2_target = copy.deepcopy(self.q2)
self.log_alpha = torch.tensor(cfg.init_alpha, device=device).log().requires_grad_(cfg.learn_alpha)
self.target_entropy = cfg.target_entropy if cfg.target_entropy is not None else -float(num_actions)
self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=cfg.learning_rate)
self.critic_optimizer = torch.optim.Adam(
list(self.q1.parameters()) + list(self.q2.parameters()), lr=cfg.learning_rate
)
self.alpha_optimizer = torch.optim.Adam([self.log_alpha], lr=cfg.learning_rate)
self.replay = ReplayBuffer(cfg.replay_capacity, env.num_envs, obs_dim, num_actions, device)
self.learning_rate = cfg.learning_rate # read by the runner for logging
self.diagnostics = {} # published as Diagnostics/<key>
self._pending = None # the obs/action awaiting its outcome
# --- collection -------------------------------------------------------------------
def get_obs_set(self, obs, set_name):
groups = self.obs_groups[set_name]
if len(groups) == 1:
return obs[groups[0]]
return torch.cat([obs[group] for group in groups], dim=-1)
def act(self, obs: TensorDict) -> torch.Tensor:
actor_obs = self.get_obs_set(obs, "actor")
actions = self.actor(actor_obs, stochastic=True).detach()
self._pending = (actor_obs, actions)
return actions
def process_env_step(self, obs, rewards, dones, extras) -> None:
actor_obs, actions = self._pending
self.actor.update_normalization(self.get_obs_set(obs, "actor"))
# A time-limit truncation is not a real terminal state, so it must not stop the bootstrap
terminal = dones.float()
if "time_outs" in extras:
terminal = (dones & ~extras["time_outs"].to(dones.device)).float()
self.replay.add(actor_obs, actions, rewards, self.get_obs_set(obs, "actor"), terminal)
def compute_returns(self, obs) -> None:
"""No-op: GAE has no meaning off-policy."""
# --- update -----------------------------------------------------------------------
def update(self) -> dict[str, float]:
if self.replay.size < self.cfg.learning_starts:
return {"critic": 0.0, "actor": 0.0, "alpha": 0.0}
alpha = self.log_alpha.exp().detach()
totals = {"critic": 0.0, "actor": 0.0, "alpha": 0.0}
for _ in range(self.cfg.gradient_steps):
obs, actions, rewards, next_obs, dones = self.replay.sample(self.cfg.batch_size)
# Critics: regress onto the entropy-regularized soft Bellman target
with torch.no_grad():
next_actions = self.actor(next_obs, stochastic=True)
next_log_prob = self.actor.get_output_log_prob(next_actions)
next_q = torch.min(
self.q1_target(torch.cat([next_obs, next_actions], dim=-1)),
self.q2_target(torch.cat([next_obs, next_actions], dim=-1)),
)
target = rewards + self.cfg.gamma * (1.0 - dones) * (next_q - alpha * next_log_prob)
q_input = torch.cat([obs, actions], dim=-1)
critic_loss = (
torch.nn.functional.mse_loss(self.q1(q_input), target)
+ torch.nn.functional.mse_loss(self.q2(q_input), target)
)
self.critic_optimizer.zero_grad()
critic_loss.backward()
torch.nn.utils.clip_grad_norm_(
list(self.q1.parameters()) + list(self.q2.parameters()), self.cfg.max_grad_norm
)
self.critic_optimizer.step()
# Actor: maximize Q minus the entropy penalty
fresh_actions = self.actor(obs, stochastic=True)
log_prob = self.actor.get_output_log_prob(fresh_actions)
q_input = torch.cat([obs, fresh_actions], dim=-1)
actor_loss = (alpha * log_prob - torch.min(self.q1(q_input), self.q2(q_input))).mean()
self.actor_optimizer.zero_grad()
actor_loss.backward()
torch.nn.utils.clip_grad_norm_(self.actor.parameters(), self.cfg.max_grad_norm)
self.actor_optimizer.step()
# Temperature: drive the policy's entropy towards the target
alpha_loss = torch.zeros((), device=self.device)
if self.cfg.learn_alpha:
alpha_loss = -(self.log_alpha.exp() * (log_prob.detach() + self.target_entropy)).mean()
self.alpha_optimizer.zero_grad()
alpha_loss.backward()
self.alpha_optimizer.step()
alpha = self.log_alpha.exp().detach()
# Polyak-average the targets
with torch.no_grad():
for online, target_net in ((self.q1, self.q1_target), (self.q2, self.q2_target)):
for p, p_target in zip(online.parameters(), target_net.parameters()):
p_target.mul_(1.0 - self.cfg.tau).add_(self.cfg.tau * p)
totals["critic"] += critic_loss.item()
totals["actor"] += actor_loss.item()
totals["alpha"] += float(alpha_loss)
self.diagnostics = {"alpha": float(alpha), "replay_size": float(self.replay.size)}
return {key: value / self.cfg.gradient_steps for key, value in totals.items()}
# --- the rest of the protocol -------------------------------------------------------
def train_mode(self):
for model in (self.actor, self.q1, self.q2):
model.train()
def eval_mode(self):
for model in (self.actor, self.q1, self.q2):
model.eval()
def get_policy(self):
"""The actor, which is what export() ships."""
return self.actor
def get_action_std(self):
return float(self.actor.output_distribution_params[1].mean()) if hasattr(
self.actor, "output_distribution_params"
) else None
def save(self) -> dict:
return {
"actor_state_dict": self.actor.state_dict(),
"q1_state_dict": self.q1.state_dict(),
"q2_state_dict": self.q2.state_dict(),
"log_alpha": self.log_alpha.detach(),
}
def load(self, loaded_dict, load_cfg, strict) -> bool:
self.actor.load_state_dict(loaded_dict["actor_state_dict"], strict=strict)
self.q1.load_state_dict(loaded_dict["q1_state_dict"], strict=strict)
self.q2.load_state_dict(loaded_dict["q2_state_dict"], strict=strict)
self.q1_target, self.q2_target = copy.deepcopy(self.q1), copy.deepcopy(self.q2)
self.log_alpha = loaded_dict["log_alpha"].clone().requires_grad_(self.cfg.learn_alpha)
return True
def compile(self, mode=None):
"""No-op: compiling is optional, and the runner only asks."""
def broadcast_parameters(self):
raise NotImplementedError("SAC here is single-device.")
def reduce_parameters(self):
raise NotImplementedError("SAC here is single-device.")Running it
cfg = OnPolicyRunnerConfig(
obs_groups={"actor": ["observation"], "critic": ["observation"]},
num_learning_iterations=2000,
# Off-policy wants short collection chunks between updates
num_steps_per_env=8,
algorithm=SACConfig(class_name=SAC, learning_rate=3e-4, batch_size=256, gradient_steps=64),
actor=MLPConfig(
hidden_dims=(256, 256),
obs_normalization=True,
# Squashed, so actions stay inside the action range SAC assumes
distribution_cfg=GaussianDistributionConfig(class_name="SquashedGaussianDistribution"),
),
critic=MLPConfig(hidden_dims=(256, 256)), # built, unused: SAC makes its own Q networks
)
runner = create_runner(env=env, runner_cfg=cfg, device="cuda:0")
runner.learn()
runner.export() # unchanged: export ships get_policy(), which is the actorTwo things worth noting about this. Deployment does not change — export() calls get_policy(), so the ONNX file is the same shape as a PPO one and Deploy a Policy applies verbatim. And critic is still built from the config and handed over, costing one unused network; the runner has no way to skip it.
What you inherit from PPO's implementation by not writing SAC — and therefore have to supply yourself above — is the same short list every from-scratch algorithm faces: observation-set selection, the normalizer updates that keep the exported policy honest, time-limit handling from extras["time_outs"], and a checkpoint dict that save(), load() and export() agree on.
Replacing the runner instead
If what you need to change is the loop rather than the update — a different collection schedule, an off-policy runner, extra bookkeeping per iteration — subclass the runner instead of the algorithm. runner.class_name is a real extension point, and a more permissive one than the algorithm's:
from telekinesis.rlbotics.config import OnPolicyRunnerConfig
from telekinesis.rlbotics.runner import OnPolicyRunner, create_runner
class MyRunner(OnPolicyRunner):
"""Same protocol, different loop."""
cfg = OnPolicyRunnerConfig(
obs_groups={"actor": ["policy"], "critic": ["policy"]},
# The class itself, a registered name, or an import path — all three work here
class_name="my_pkg.runner:MyRunner",
)
runner = create_runner(env=env, runner_cfg=cfg, device="cuda:0")create_runner tries RUNNERS first — which holds {"OnPolicyRunner": OnPolicyRunner} and is meant to grow — and falls back to importing a "module:Class" path, so a custom runner has to be importable but not registered. That is the one place where the string form works, and the reason it works there and not for the algorithm.
create_runner, not the class
OnPolicyRunner(...) always builds that one class, so class_name is silently ignored when you construct the runner directly. Only create_runner reads it.
Things that are not extension points
| Thing | Status |
|---|---|
rnd_cfg (Random Network Distillation) | Configured for but not implemented — setting it raises. Do not treat it as a hook |
| Symmetry with a recurrent policy | Rejected at construction |
CUDA-graph torch.compile modes | Rejected: incompatible with the multi-model forward pattern the algorithms use |
A non-PPOConfig algorithm config | Rejected with Unsupported algorithm config type. Subclass PPOConfig rather than writing a config from scratch |