TBM

Trajectory Bundle Method (TBM) solver for constrained trajectory optimization.

This repository provides:

• tbm/: the installable package
• examples/: runnable toy problems
• tests/: unit and integration tests

Install

For development and tests:

uv sync

Optional runtime extras:

uv sync --extra cvxpy
uv sync --extra mjwarp --extra examples

The default subproblem backend is clarabel_direct, so the core API does not require cvxpy. Install the cvxpy extra only if you want the alternate CVXPY-backed subproblem path.

Documentation

The public API guide lives in docs/public_api.md. It covers install options, callback shapes, solver configuration, and the result object returned by solve().

Quick Start

import numpy as np

from tbm import InitialGuess, TBMConfig, TBMProblem, solve


def rollout(batch_x0, batch_u):
    dt = 0.1
    states = np.zeros((batch_x0.shape[0], batch_u.shape[1] + 1, 2))
    states[:, 0] = batch_x0
    for t in range(batch_u.shape[1]):
        x = states[:, t]
        u = batch_u[:, t]
        states[:, t + 1, 0] = x[:, 0] + dt * x[:, 1]
        states[:, t + 1, 1] = x[:, 1] + dt * u[:, 0]
    return states


problem = TBMProblem(
    state_dim=2,
    control_dim=1,
    horizon=16,
    initial_state=np.array([0.0, 0.0]),
    rollout=rollout,
    terminal_cost_residual=lambda x: np.column_stack(
        [
            np.sqrt(10.0) * (x[:, 0] - 1.0),
            x[:, 1],
        ],
    ),
)

guess = InitialGuess(controls=np.zeros((15, 1)))
result = solve(problem, guess, TBMConfig(max_iterations=20))
print(result.status, result.max_constraint_violation)

For solver logging, use TBMConfig(verbose=True) to print outer TBM iteration statistics. Use TBMConfig(cvxpy_verbose=True) only when you want the inner CVXPY or Clarabel solver output. Select TBMConfig(subproblem_backend="cvxpy") only when the cvxpy extra is installed.

Examples

• uv run python examples/double_integrator.py
• uv run --extra mjwarp --extra examples python examples/cartpole_mjwarp.py
• uv run --extra cvxpy --extra examples python examples/quadcopter.py
• uv run --extra cvxpy --extra examples python -m examples.race_car.sqp
• uv run --extra cvxpy --extra examples python -m examples.race_car.mppi

Tests

Run the default CPU-safe test suite:

uv run pytest

Run GPU integration tests:

uv run --extra mjwarp pytest -m gpu

CI

GitHub Actions runs Ruff, the default pytest suite, and a wheel smoke test from .github/workflows/ci.yml.