Deploy a Policy

SUMMARY

runner.export() writes one self-contained policy.onnx: the observation normalization the policy trained with, the deterministic action, and — where the environment has action bounds — the scaling onto them, all inside the graph.

Running it needs neither PyTorch nor RLBotics. Only numpy and onnxruntime.

Export

policy_path = runner.export()
# logs/gymnasium_ppo/2026-08-06_18-08-47/policy.onnx

Argument	Default	What it does
path	None	Directory to write to. None means this run's directory
filename	"policy.onnx"	Name of the exported file
from_best	True	Export the run's best checkpoint. False exports the policy currently in memory

It exports the best checkpoint, not the last one

This is the detail that matters most. export() ships the run's best checkpoint, not whichever policy training happened to end on. Those differ whenever the reward peaked and then fell back, which for a long locomotion run is the normal case — shipping the final policy there means deploying a worse one than you trained.

model_best.pt is rewritten whenever the mean episode reward beats every earlier iteration's, and it is checked every iteration rather than on the save_interval grid, so a peak the reward later falls back from is not lost.

Use runner.export(from_best=False) when you specifically want to compare where training ended up against the best it found.

WARNING

export() needs somewhere to write. With logging off and no path given it raises — pass a path, or set logger.log_dir.

Run it

from telekinesis.rlbotics.policy import Policy

policy = Policy("logs/gymnasium_ppo/2026-08-06_18-08-47/policy.onnx")
action = policy.get_action(observation)   # (obs_dim,) -> (num_actions,), or batched

That is the whole deployment API:

Member	Returns
Policy(path)	Loads the file. Raises FileNotFoundError if it is missing, ImportError if onnxruntime is not installed
policy.obs_dim	Number of observation values the policy expects
policy.num_actions	Number of action values it produces
policy.get_action(obs)	The deterministic action, already scaled to the environment's action range

The batch dimension is dynamic, so a single observation of shape (obs_dim,) and a batch of shape (batch, obs_dim) both work — the return shape follows the input.

The action is deterministic: the exported graph carries the mean of the trained policy rather than a sample from it, which is what deployment wants.

On a real robot

The loop is yours; the policy is a pure function inside it. What you have to get right is assembling the observation in exactly the order the environment produced it during training:

import numpy as np
from telekinesis.rlbotics.policy import Policy

policy = Policy("policy.onnx")

while True:
    # Same layout, same order, same units as the training environment
    observation = np.concatenate([
        base_lin_vel,        # 3
        base_ang_vel,        # 3
        projected_gravity,   # 3
        joint_positions,     # n
        joint_velocities,    # n
        previous_action,     # n
        velocity_command,    # 3
    ]).astype(np.float32)

    action = policy.get_action(observation)
    robot.set_joint_targets(action)

You do not need to normalize the observation — the running statistics the policy trained with are baked into the graph, so raw values go in.

A wrong observation layout fails silently

The graph takes a flat vector of obs_dim values. Swap two blocks and it still runs, still returns plausible-looking actions, and the robot behaves badly for no visible reason. Print policy.obs_dim and check it against the training environment's observation_spec() before trusting the loop.

Acting inside a GPU simulator

Deploying onto a robot needs nothing but numpy. But if you are acting in mjlab or Isaac Lab to evaluate a policy, wrap the loop in torch.inference_mode() — those simulators allocate buffers during training that cannot be updated in place outside it:

import torch

with torch.inference_mode():
    obs = env.reset()
    for _ in range(num_steps):
        actor_obs = obs[obs_groups[0]]                     # or concatenate several groups
        action = policy.get_action(actor_obs.cpu().numpy())
        obs, rewards, dones, _ = env.step(torch.as_tensor(action, device=env.device))

Torch appears here only because a vectorized environment speaks tensors — the policy itself still takes numpy in and gives numpy out.

Action scaling

Whether the exported graph scales its actions depends on the environment it was trained in:

Backend	Action bounds	What the export carries
Gymnasium	Read from the task spec	Observation normalization and clipping plus affine scaling onto the task's bounds
mjlab	None — actions are joint targets	Observation normalization only
Isaac Lab	None — actions are joint targets	Observation normalization only

So for mjlab and Isaac Lab, apply the same clip_actions limit you trained with before the actions reach the robot. For Gymnasium tasks the policy already emits actions the environment accepts directly.

The scaling, where present, maps the policy's [-1, 1] output onto [low, high]:

action = low + 0.5 * (clamp(policy(obs), -1, 1) + 1) * (high - low)

TorchScript instead of ONNX

If your deployment target already has PyTorch, export_policy_to_jit() traces the same deterministic policy to TorchScript:

runner.export_policy_to_jit("deploy", filename="policy.pt")

Note the difference: unlike export(), this traces the current policy rather than the best checkpoint, and it does not fold in the action scaling.

Not happy with the policy you exported?

What each curve means, the measured hyperparameter recipes, and why exploration collapses without a std floor.

Tuning Best Practices →