Architecture

SUMMARY

Written for someone about to open a pull request. Four responsibilities are kept apart on purpose:

The runner owns the loop. The algorithm owns the update. An adapter owns the simulator. One exported file owns deployment.

Each is replaceable without touching the others, which is why a custom environment, algorithm or runner is a small job rather than a fork. Source: telekinesis-ai/telekinesis-rlbotics.

Architecture Overview

config.yamlOne file describes the whole run

env:runner:

↓

VecEnv adapterThe simulator, behind one interface

GymnasiumVecEnvMjlabVecEnvIsaacLabVecEnv

↓

create_runner → OnPolicyRunnerBuilds everything the config describes, then owns the loop

actor + criticPPORolloutBufferLoggerCheckpointManager

↓

learn()collect → bootstrap → update → log → checkpoint

logs/<experiment>/<timestamp>/

↓

export() → policy.onnxSelf-contained: numpy in, actions out

Policyyour robot

A config names a simulator and a training setup. create_runner builds everything that setup describes. learn() turns environment steps into gradient steps. export() reduces the result to one file that needs neither PyTorch nor RLBotics to run.

The loop, in eight lines

This is the whole of learn() with the bookkeeping removed. Everything else in the library exists to serve these calls, so it is the fastest way to orient yourself in runner.py:

for iteration in range(num_learning_iterations):
    for _ in range(num_steps_per_env):          # collect
        actions = alg.act(obs)
        obs, rewards, dones, extras = env.step(actions)
        alg.process_env_step(obs, rewards, dones, extras)

    alg.compute_returns(obs)                    # bootstrap
    losses = alg.update()                       # learn
    logger.log(...)                             # report
    checkpoints.save_best(...)                  # keep the best policy

Collection runs inside torch.inference_mode(), update() does not

Tensors created during collection cannot be saved for backward, which is why RolloutBuffer pre-allocates its storage and copies into it rather than appending. Anything you write that buffers transitions has to do the same.

The loop is on-policy by construction

Collect a fixed chunk, update once, clear the storage. An off-policy algorithm fits by keeping its own replay buffer and treating the chunk as a collection cadence — see Introduce a Custom Algorithm.

Folder structure

telekinesis-rlbotics/
├── src/telekinesis/rlbotics/   the library — nothing here imports a simulator at module level
│   └── envs/                   the VecEnv interface and one adapter per simulator
├── configs/                    one YAML per task, grouped by framework, plus example.yaml
│   ├── gymnasium/              12 tasks
│   ├── mjlab/                  10 tasks
│   └── isaaclab/               31 tasks
├── examples/                   runnable entry points
│   └── module_examples/        one script per library piece, no simulator needed
├── tests/                      pytest, one file per module
├── scripts/windows/            smoke_test.bat, runs training_example.py over every config
├── .github/workflows/          verify.yml, develop.yml, release.yml
└── pyproject.toml              deps, the four extras, pytest config

Directory	Why you would open it
src/telekinesis/rlbotics/	The library. Every module is listed below
src/telekinesis/rlbotics/envs/	Adding a simulator, or changing how an existing one is presented
configs/	Adding a task, or tuning one. Start from example.yaml
examples/	The public entry points. Any API change should keep these working
tests/	One file per module — a new module needs a new test_*.py

Note the import discipline in envs/: each adapter imports its simulator lazily, inside a try, so installing one extra never requires the others. Preserve that when touching an adapter.

Module by module

Core

Module	Lines	What it contains
runner.py	774	OnPolicyRunner — builds the models, storage, algorithm, logger and checkpoint manager from a config, then owns the training loop, the ONNX/JIT export and the multi-GPU setup. Plus create_runner and the RUNNERS registry
config.py	773	Every typed config dataclass, each validating itself in __post_init__: OnPolicyRunnerConfig, PPOConfig, LoggerConfig, MLPConfig, CNNConfig, CNNEncoderConfig, GaussianDistributionConfig, SymmetryConfig, and the BaseConfig that gives them all to_dict / from_dict / replace
algorithms.py	564	PPO — the surrogate and value losses, the KL-adaptive learning rate, the mini-batch loop. Plus the ALGORITHMS and OPTIMIZERS registries, check_nan and compile_model
rollout.py	416	Transition, RolloutBuffer and RolloutBufferBatch — one rollout's worth of pre-allocated storage, GAE, and the mini-batch generator

Models

Module	Lines	What it contains
models.py	765	MLPModel and CNNModel, the MLP and CNN building blocks, build_distribution and the DISTRIBUTIONS registry. Also the export-friendly deterministic forms as_jit() / as_onnx()
distributions.py	304	Distribution, GaussianDistribution and SquashedGaussianDistribution — what makes an actor stochastic, and the std parameterization and clamping
normalization.py	167	EmpiricalNormalization and EmpiricalDiscountedVariationNormalization — running observation statistics, which travel inside the exported policy

Environments

Module	Lines	What it contains
envs/base.py	101	The VecEnv interface — the whole contract between the library and any simulator — plus observation_spec
envs/gym_env.py	247	GymnasiumVecEnv. numpy↔torch conversion, and mapping [-1, 1] onto the task's action bounds
envs/mjlab_env.py	313	MjlabVecEnv. Reads the observation groups, action dimension and episode limit off a manager-based task; forwards truncations only on an infinite-horizon task
envs/isaaclab_env.py	371	IsaacLabVecEnv, plus launch_simulator / shutdown_simulator — Isaac Sim has to be running before any task is imported, and that ordering lives here

Run output

Module	Lines	What it contains
logger.py	553	Logger — TensorBoard scalars, the console report, the run directory layout — and VideoLogger, which renders from the training rollout itself
checkpoint.py	245	CheckpointManager — the save interval, rotation, model_best.pt scored every iteration, and the resume lookup that searches every run of an experiment
policy.py	87	Policy. The entire deployment surface: load an ONNX file, numpy in, actions out. Imports neither torch nor anything else in the library

Extensions and helpers

Module	Lines	What it contains
symmetry.py	130	Symmetry — mirrored mini-batch augmentation and the mirror loss. Separate because the mirror function comes from outside the library
utils.py	109	set_seed, resolve_device (including the MPS-vs-CPU note) and resolve_callable, which imports a "module:Class" path

That is roughly 5,500 lines. runner.py and config.py are the two big ones and the two that most changes touch.

The four seams

Where to extend, and how each is selected:

Seam	What you write	Selected by	Guide
Environment	A VecEnv, or an mjlab / Isaac Lab task	Passed as an object, not a name	Custom Environment
Algorithm	A class implementing the algorithm protocol	algorithm.class_name, as the class object	Custom Algorithm
Runner	A subclass of OnPolicyRunner	runner.class_name, by name or import path	Replacing the runner
Model / distribution	A model or distribution class	actor.class_name, distribution_cfg.class_name	Configuration

Only the runner's class_name accepts an import path

create_runner falls back to resolve_callable for an unregistered name, so a custom runner has to be importable but not registered. The algorithm, models and distributions resolve through fixed registries — ALGORITHMS, DISTRIBUTIONS and the model equivalent — so pass the class object there instead of a string.

One design consequence worth knowing before you touch the export path: only the actor is exported. Anything the critic reads is free at deployment, which is what makes an asymmetric actor-critic a config change rather than a code change; anything the actor reads becomes a permanent obligation on the robot.

Development

git clone https://github.com/telekinesis-ai/telekinesis-rlbotics.git
cd telekinesis-rlbotics
pip install -e ".[dev,gym]"

pytest                                   # tests/, one file per module
ruff check .
pylint src
python examples/run_all_examples.py      # smoke-tests everything installed, SKIPs the rest

CI lives in .github/workflows/ — verify.yml on pull requests, develop.yml, and release.yml for publishing to PyPI. ruff and pylint are pinned in the dev extra, so match their versions rather than your local ones.

Two conventions the codebase holds to consistently, worth matching in a pull request: every config validates in __post_init__ rather than trusting its caller, and docstrings explain the why — the measured numbers behind a default, or the failure a check prevents — because that is the context a reader cannot recover from the code.

Every setting, with its type and default

The annotated reference YAML, the config tables, and how to build the same configuration in Python.

Configuration →