Optimization
Elements offset to offset + batch_size, sharing storage.
Move by velocity; only shortened steps write out_velocity, which may alias the input velocity.
(batch, tangent_dim) velocities that integrate moves each element by to reach other's.
Accept proposed values and derived state where accept_mask[i] == 1; leave other rows unchanged.
Combine Vars into tangent-column blocks; preserve subclass step rules across copies and views.
A batched vector optimization variable on Euclidean space.
Example
>>> from enum import Enum >>> from robokit.opt import VarValues >>> class Key(Enum): ... X = 0 >>> var = VectorVar(wp.array([[1.0, 2.0]], dtype=wp.float32)) >>> values = VarValues.from_var(var, Key.X) >>> result = values.integrate(wp.array([[0.5, -1.0]], dtype=wp.float32)) >>> result.get(Key.X).x.numpy().tolist() [[1.5, 1.0]] >>> capped = VectorVar(wp.zeros((1, 2), dtype=wp.float32), unit=10.0, max_step=1.0) # moves 10x, 1 at most >>> capped.integrate(wp.array([[0.0625, 1.0]], dtype=wp.float32)).x.numpy().tolist() [[0.625, 1.0]]
import warp as wp
from robokit.opt.var_values import VectorVar
from enum import Enum
from robokit.opt import VarValues
class Key(Enum):
X = 0
var = VectorVar(wp.array([[1.0, 2.0]], dtype=wp.float32))
values = VarValues.from_var(var, Key.X)
result = values.integrate(wp.array([[0.5, -1.0]], dtype=wp.float32))
result.get(Key.X).x.numpy().tolist()
capped = VectorVar(wp.zeros((1, 2), dtype=wp.float32), unit=10.0, max_step=1.0) # moves 10x, 1 at most
capped.integrate(wp.array([[0.0625, 1.0]], dtype=wp.float32)).x.numpy().tolist()Warp optimizer lifecycle with shared Jacobian machinery for gradient optimizers and LM's CG step.
Solve sequence:
reset: build once, slice to batch size, optionally capture a graph, then call_start.- Repeat
stepmax_itertimes, replaying the iteration graph or running_step_body.
Subclass hooks:
_build(var): allocate all buffers at the first var's batch size (max_batch_size)._BATCH_BUFFERS: identify batch-leading arrays and vars to slice per solve._start(var): initialize per-solve damping/history/schedules and starting cost._step_body(var): run one iteration in place using kernels only.
Optimize var in place; returns it with the final cost per batch element.
Prepare to iterate from var; num_steps (default max_iter) steps follow and L-BFGS finishes the last.
One iteration in place; an outer capture records kernels instead of replaying a graph.
Iterations per solve.
Check the stop rule every this many iterations; 0 never stops early (each check syncs the host).
The solve stops once every batch element's cost is at most this.
Capture per width: "iter" replays one step; "full" replays the solve without callbacks or early stopping.
The optimizer a PopulationSolver steps with this config.
First-order optimizer supporting GD and Adam, with optional decoupled weight decay (AdamW).
Optimize var in place; returns it with the final cost per batch element.
Prepare to iterate from var; num_steps (default max_iter) steps follow and L-BFGS finishes the last.
Iterations per solve.
Check the stop rule every this many iterations; 0 never stops early (each check syncs the host).
The solve stops once every batch element's cost is at most this.
Capture per width: "iter" replays one step; "full" replays the solve without callbacks or early stopping.
Step size.
Learning rate per iteration from (zero-based iteration, learning_rate); None keeps it constant.
analytic_jacobian: terms fill a shared sparse/dense J, then one Jᵀr reduction; autodiff: tape gradients; analytic_gradient: term kernels accumulate cost and gradient directly.
Update rule: plain gradient descent or Adam.
Clamp each gradient entry to ±this before the update, so a stiff cost cannot overshoot.
Adam's first-moment decay.
Adam's second-moment decay.
Adam's denominator stabilizer.
Decoupled decay: each step also moves x by -learning_rate weight_decay x (AdamW under Adam).
GD with RMSProp proposals and Metropolis acceptance, supporting every base-class gradient mode.
Acceptance outcomes:
- Accept: keep the proposed point and gradient.
- Reject: restore the previous point and gradient, continuing the chain there.
Optimize var in place; returns it with the final cost per batch element.
Prepare to iterate from var; num_steps (default max_iter) steps follow and L-BFGS finishes the last.
learning_rate and the Metropolis temperature both decay by temperature_decay, on their own periods.
Iterations per solve.
Check the stop rule every this many iterations; 0 never stops early (each check syncs the host).
The solve stops once every batch element's cost is at most this.
Capture per width: "iter" replays one step; "full" replays the solve without callbacks or early stopping.
The optimizer a PopulationSolver steps with this config.
Initial RMSProp step size.
How the gradient is computed, as in GDOptimizerConfig.
RMSProp's weight on the newest squared gradient.
Iterations between step-size decays.
Initial Metropolis temperature.
Factor applied to the step size and the temperature at each of their periods.
Iterations between temperature decays.
RMSProp's denominator stabilizer.
Seed of the Metropolis acceptance draws.
Optimize var in place; returns it with the final cost per batch element.
Prepare to iterate from var; num_steps (default max_iter) steps follow and L-BFGS finishes the last.
One iteration in place; an outer capture records kernels instead of replaying a graph.
Iterations per solve.
The solve stops once every batch element's cost is at most this.
Capture per width: "iter" replays one step; "full" replays the solve without callbacks or early stopping.
Off by default: without host checks, a folded line search defers each step into the next direction launch.
Curvature pairs kept for the two-loop recursion.
Initial inverse-Hessian scale when there is no history or normalize_direction is on (else Barzilai-Borwein).
Trial step sizes, evaluated in parallel.
Sufficient-decrease constant.
Curvature constant (approx_wolfe).
armijo: largest passing trial, else no step; approx_wolfe: largest Armijo/curvature trial, else smallest, and needs folded gradients.
Normalize directions so trial steps are lengths; H0 then uses h0_scale instead of Barzilai-Borwein.
analytic_jacobian: terms fill a shared sparse/dense J, then one Jᵀr reduction; autodiff: tape gradients; analytic_gradient: term kernels accumulate cost and gradient directly.
Fold trial gradients to skip a recompute; needs analytic_gradient or all-sparse analytic_jacobian terms.
Runs without a per-iteration graph: each step branches on the host iteration count.
Levenberg-Marquardt with Cholesky or conjugate gradients.
Solver paths:
- Cholesky/default Jacobian terms: fused fast path.
- Cholesky/other terms: accumulate
JᵀJ, then factor in one tile or 16×16 tiles within the sparse-pattern band. - CG: reuse sparse/dense Jacobians without forming
JᵀJ.
Optimize var in place; returns it with the final cost per batch element.
Prepare to iterate from var; num_steps (default max_iter) steps follow and L-BFGS finishes the last.
One iteration in place; an outer capture records kernels instead of replaying a graph.
Iterations per solve.
Check the stop rule every this many iterations; 0 never stops early (each check syncs the host).
The solve stops once every batch element's cost is at most this.
Capture per width: "iter" replays one step; "full" replays the solve without callbacks or early stopping.
Initial damping.
Damping divides by this on accept and multiplies on reject.
Smallest damping.
Largest damping.
Minimum gain ratio rho to accept a step.
Stabilizer in the rho denominator; 0 uses the exact ratio.
Native Cholesky, matrix-free CG, or optional Torch float64 Cholesky for large dense systems.
CG iterations per step; None: max(4, min(64, tangent_dim // 4)).
Squared CG residual that stops early; checked on the host, so only without CUDA graphs.
CG iterations between those host checks.
CG and Torch Cholesky use eager steps; the native backend retains CUDA graph capture.
LMOptimizerConfig(linear_solver="cholesky"), the default, solves each step exactly. Small systems factor in one tile; larger ones use the band implied by SparseTask sparsity patterns. linear_solver="cg" iterates matrix-free on whichever Jacobian the terms provide, sparse or dense, and never forms the normal matrix.
Refine batch * num_candidates[0] VarValues, reusing one var, allocation, and graph per population width.
Each scheduled round:
sample: move candidates using prior scores, then draw around them.optimize: move each candidate toward lower cost.score: evaluate the terms or score term.select: keep each problem's bestnum_candidates[k + 1]candidates.
(batch, 1) row of each instance's winner: the first ranked index.
Run the schedule on widest var; return each instance's best candidate and cost.
Sum final-segment term costs by cost_name, including duplicate names such as ConfigurationTask.
What a round does, and when: stages run in order, iterations counted across them.
The schedule, e.g. [Stage(6, 64, LMOptimizerConfig(lm_lambda=10.0)), Stage(10, 4), Stage(0, 1)]: 64 candidates for 6 iterations, the best 4 for 10 more, then the best one.
Built during setup, then applied in order where the stages sample.
At each sampling iteration:
- Sample the population before optimizing.
- Restart the optimizer on that population.
- At the last iteration, sample after the final score and end on the moved candidates.
"iter" replays one optimizer step per width; "full" the whole solve (no callbacks or early stopping).
Let a segment end at its optimizer's stop rule (OptimizerConfig.early_stopping_interval); not under "full".
Score only the final population and return it best first; earlier selects keep the optimization cost.
If positive, a drawing sampler's solve ends once each instance's cheapest draw beats its first one (the center) by less than this, patience scores in a row; the last sample still moves the candidates. Reads scores on the host.
The widest the population gets per candidate: the drawing sampler's draw count, or 1.
Iterations of the whole schedule.
Iterations (1-based, after the update) where a stage ends and the next one's width is kept.
Population width per instance in each stage.
Each stage's gradient update, or None if no stage sets one.
Zero-based sampling iterations; empty: the default (see Stage.sample_every).
A run of iters iterations at one population width. Entering a stage keeps each instance's best num_candidates (the first stage's is the input width); unset fields keep the previous stage's.
Iterations of this stage.
Population width per instance; None: the previous stage's, or 1.
Gradient update; None: the previous stage's. With none in any stage: drawing samplers (MPPI/ES) update by sampling, other solvers by dense LM.
Draw with the samplers every this many iterations from the stage's start; 0 in every stage: a drawing sampler draws before every score and once after the last, others at each select.
Sampler settings; owners supply variables, population width, device and masks at build time.
Construct an independent runtime sampler for these variables.
Low-discrepancy initialization settings; helpers may scale the spread by joint ranges.
Example
>>> config = RobertsSamplerConfig(scale=0.5, keep_first=True) >>> config.scale, config.keep_first, config.num_draws (0.5, True, 1)
from robokit.opt.sampler import RobertsSamplerConfig config = RobertsSamplerConfig(scale=0.5, keep_first=True) config.scale, config.keep_first, config.num_draws
Gaussian draw/update settings, independent of runtime population buffers.
Example
>>> from robokit.opt.var_values import VectorVar
>>> class Key(Enum):
... X = 0
>>> var = VarValues({Key.X: VectorVar(wp.zeros((1, 2), dtype=wp.float32, device="cpu"))})
>>> config = GaussianSamplerConfig(var_key=Key.X, num_draws=8)
>>> first, second = config.build(var, 1), config.build(var, 1)
>>> first is second
False
>>> first.noise.shape
(8, 2)import warp as wp
from enum import Enum
from robokit.opt.sampler import GaussianSamplerConfig
from robokit.opt.var_values import VarValues
from robokit.opt.var_values import VectorVar
class Key(Enum):
X = 0
var = VarValues({Key.X: VectorVar(wp.zeros((1, 2), dtype=wp.float32, device="cpu"))})
config = GaussianSamplerConfig(var_key=Key.X, num_draws=8)
first, second = config.build(var, 1), config.build(var, 1)
first is second
first.noise.shapeCandidates drawn around each center, the center itself first.
How scored draws become weights: MPPI's softmax, or evolution strategies' fitness shaping.
Noise standard deviation at the start of a solve.
Floor the softmax refit keeps the standard deviation above.
Softmax temperature; lower follows the cheapest draws more closely.
Share of the weighted spread the softmax refit moves the standard deviation toward.
Fitness-shaping step size.
Guard against division by a zero spread.
Noise seed; the same seed replays the same draws.
Move each center to its draws' weighted mean, or onto its cheapest draw.
If positive, the softmax temperature is each center's cost spread over this best share of its draws (at least two) instead of temperature, and best_draws keeps its cheapest draw.
With elite_ratio, move each center to the weighted mean of that share alone.
Each scored round scales std by this before the refit, never below min_std.
Share of a solve's first-round draws (its last ones) drawn at wide_scale times std, to explore.
Draws new candidates around each candidate of one named variable.
Candidates drawn around each one; 1 draws in place.
Reopen the distribution at the start of a solve.
Move candidates by scores; return var or sampler-owned draws for scoring.
Seed candidates once with a Roberts low-discrepancy sequence, covering a box more evenly than random draws.
Sampling options:
scale: half-width per tangent column; joint columns usually use joint range.sampler.scale.assign(...): change the device-side scale between solves.keep_first: leave candidate zero at the seed by assigning it zero noise.
Example
>>> sampler = RobertsSampler(num_candidates=4, scale=[0.5, 0.5], device="cpu") >>> tangent = wp.zeros((8, 2), dtype=wp.float32, device="cpu") >>> sampler.fill(tangent) >>> tangent.numpy().shape (8, 2) >>> bool(np.abs(tangent.numpy()).max() <= 0.25) True
import numpy as np import warp as wp from robokit.opt.sampler import RobertsSampler sampler = RobertsSampler(num_candidates=4, scale=[0.5, 0.5], device="cpu") tangent = wp.zeros((8, 2), dtype=wp.float32, device="cpu") sampler.fill(tangent) tangent.numpy().shape bool(np.abs(tangent.numpy()).max() <= 0.25)
A Roberts spread is fixed.
A Roberts spread is fixed, so scores tell it nothing: spread var in place.
Fit a sampling distribution from costs, as in MPPI and evolution strategies.
Each round:
- Move centers using scored draws and refit
std. - Draw
num_drawsstd-scaled normals per center; draw zero has no noise and preserves the center.
Center updates:
sample_mode=mean: cost-weighted mean of noise.sample_mode=best: cheapest draw's current position, including any optimizer movement.
Cost weighting:
- softmax: Boltzmann weights over negative cost at
temperature. - fitness_shaping: standardized costs scaled by
learning_rate.
Capacity/reset controls:
instances: maximum center population; narrower solves use prefixes.std_resetandtemperatures: caller-adjustable per-instance settings;resetrestores std fromstd_reset.lowerandupper: optional(instances, tangent_dim)bounds on a vector variable's drawn values; each draw is clamped to its center's and its noise follows.
Reopen the spread to std_reset and replay the same noise sequence.