Skip to content

CLAUDE.md

Latest commit

 

History

History
65 lines (46 loc) · 2.75 KB

CLAUDE.md

File metadata and controls

65 lines (46 loc) · 2.75 KB
Raw

CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

Project Overview

TBM (Trajectory Bundle Method) is a solver for constrained trajectory optimization using Sequential Convex Programming with bundle sampling. The solver is backend-agnostic; MJWarp (MuJoCo on GPU) is the first rollout backend.

Commands

# Install (core solver only)
uv sync

# Install with GPU backend and examples
uv sync --extra mjwarp --extra examples

# Run all fast (CPU) tests
uv run pytest

# Run a single test file or test
uv run pytest tests/test_solver_core.py
uv run pytest tests/test_solver_core.py::test_name

# Run GPU tests (requires NVIDIA GPU + mujoco-warp)
uv run --extra mjwarp pytest -m gpu

# Lint
uv run ruff check tbm/ tests/ examples/

# Run examples
uv run python examples/double_integrator.py
uv run --extra mjwarp --extra examples python examples/cartpole_mjwarp.py

Architecture

The solver loop lives in tbm/solver.py (solve() function):

  1. Sample bundles around the current trajectory iterate (_build_segment_bundle)
  2. Solve a convex subproblem via CVXPY+Clarabel (_solve_subproblem)
  3. Recover a new iterate from simplex-weighted samples (_recover_iterate)
  4. Evaluate feasibility & optimality on the nonlinear problem (_evaluate_iterate)
  5. Adapt trust regions and check convergence

Key modules:

  • tbm/problem.py — TBMProblem (problem definition), InitialGuess, RolloutBackend protocol, cost residual and constraint callback signatures
  • tbm/solver.py — TBMConfig (hyperparameters), SegmentBundle, TBMResult, the main SCP loop and convex subproblem assembly
  • tbm/sampling.py — Bundle sampling strategies (coordinate, gaussian, uniform) with trust-region scaling
  • tbm/mjwarp_backend.py — GPU-accelerated batched rollouts via MuJoCo-Warp

Core design decisions:

  • One simplex variable alpha per shooting segment; the convex subproblem interpolates sampled rollout data using these weights.
  • Costs are expressed as residual vectors (squared norms form the objective), not scalar costs.
  • Dynamics matching and constraints use L1-penalized slack variables.
  • segment_length controls shooting granularity: 1 = full multiple shooting, horizon-1 = single segment.
  • The RolloutBackend protocol (rollout(initial_states, control_segments, start_index) -> states) allows swapping in any batched dynamics simulator.

Testing

Tests are in tests/. GPU tests are marked with @pytest.mark.gpu and @pytest.mark.slow. CPU tests run by default; GPU tests require pytest -m gpu with the mjwarp extra installed.

Python & Tooling

  • Python >= 3.11, build system is hatchling, dependency management via uv
  • Ruff for linting (line-length 120, rules: E, F, I, B)