Terms
Base of every cost term: names its variable (var_key) and cost row (cost_name).
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Label a config gave this term, so config.task(name) finds it again.
The robot a built task reads; None on a template.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Swap each Context placeholder, alone or in a tuple, for its value; subclasses then build for the variable.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> task = Step(dt=(Inputs.STEP, 2.0))
>>> task.resolve(None, {Inputs.STEP: 0.5}) # the base swaps placeholders only
>>> task.dt
(0.5, 2.0)from robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
task = Step(dt=(Inputs.STEP, 2.0))
task.resolve(None, {Inputs.STEP: 0.5}) # the base swaps placeholders only
task.dtVariables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare shared state before this term is evaluated.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Term consumed as a weighted residual and optional Jacobian.
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Label a config gave this term, so config.task(name) finds it again.
The robot a built task reads; None on a template.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Swap each Context placeholder, alone or in a tuple, for its value; subclasses then build for the variable.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> task = Step(dt=(Inputs.STEP, 2.0))
>>> task.resolve(None, {Inputs.STEP: 0.5}) # the base swaps placeholders only
>>> task.dt
(0.5, 2.0)from robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
task = Step(dt=(Inputs.STEP, 2.0))
task.resolve(None, {Inputs.STEP: 0.5}) # the base swaps placeholders only
task.dtVariables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare shared state before this term is evaluated.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Device of the task's Warp buffers; None before init_buffers.
Allocate Warp buffers on device; a task without any keeps none.
Dimension of the residual vector.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Per-batch scalar cost and gradient, used by GDOptimizer/LBFGSOptimizer in analytic-gradient mode.
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Label a config gave this term, so config.task(name) finds it again.
The robot a built task reads; None on a template.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Swap each Context placeholder, alone or in a tuple, for its value; subclasses then build for the variable.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> task = Step(dt=(Inputs.STEP, 2.0))
>>> task.resolve(None, {Inputs.STEP: 0.5}) # the base swaps placeholders only
>>> task.dt
(0.5, 2.0)from robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
task = Step(dt=(Inputs.STEP, 2.0))
task.resolve(None, {Inputs.STEP: 0.5}) # the base swaps placeholders only
task.dtVariables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare shared state before this term is evaluated.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Accumulate this term's weighted cost and optional direct gradient.
Optimization variables.
Shared per-batch cost buffer.
Optional shared gradient buffer. If omitted, only cost is accumulated.
Residual task with an explicitly stored sparse Jacobian.
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Label a config gave this term, so config.task(name) finds it again.
The robot a built task reads; None on a template.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Swap each Context placeholder, alone or in a tuple, for its value; subclasses then build for the variable.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> task = Step(dt=(Inputs.STEP, 2.0))
>>> task.resolve(None, {Inputs.STEP: 0.5}) # the base swaps placeholders only
>>> task.dt
(0.5, 2.0)from robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
task = Step(dt=(Inputs.STEP, 2.0))
task.resolve(None, {Inputs.STEP: 0.5}) # the base swaps placeholders only
task.dtVariables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare shared state before this term is evaluated.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Device of the task's Warp buffers; None before init_buffers.
Allocate Warp buffers on device; a task without any keeps none.
Dimension of the residual vector.
Sparsity pattern of the Jacobian matrix.
Compute the non-zero values of the weighted Jacobian matrix according to the sparsity pattern.
Accumulate normal equations directly from existing sparse rows.
Write residuals and optional sparse Jacobian values into shared buffers.
Optimization variables.
Shared residual buffer.
First residual row for this task.
Optional shared sparse-value buffer.
First sparse entry for this task.
Set the valid trajectory length for each batch element.
Per-batch valid frame counts.
Task with per-iteration work outside a captured CUDA graph.
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Label a config gave this term, so config.task(name) finds it again.
The robot a built task reads; None on a template.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Swap each Context placeholder, alone or in a tuple, for its value; subclasses then build for the variable.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> task = Step(dt=(Inputs.STEP, 2.0))
>>> task.resolve(None, {Inputs.STEP: 0.5}) # the base swaps placeholders only
>>> task.dt
(0.5, 2.0)from robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
task = Step(dt=(Inputs.STEP, 2.0))
task.resolve(None, {Inputs.STEP: 0.5}) # the base swaps placeholders only
task.dtVariables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare shared state before this term is evaluated.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Device of the task's Warp buffers; None before init_buffers.
Allocate Warp buffers on device; a task without any keeps none.
Dimension of the residual vector.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Prepare eager state before a Warp reduction.
Current or proposed optimization variables.
Whether this is a line-search proposal.
Solver iteration.
Update eager state after accepting or rejecting a step.
Per-batch accepted-step mask.
Solver iteration.
Per-batch costs.
Whether the solver should stop early.
Serial preparation shared by terms that consume a robot state.
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Label a config gave this term, so config.task(name) finds it again.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Swap placeholders, then take the robot of the state this task reads.
Prepare kinematics and geometry required by the task.
Apply independent per-pose tasks to trajectory frames, with one dense tangent block per selected frame.
Per-pose residual task.
Number of trajectory frames, or omitted so the trajectory optimizer sets it.
Frames to include. Defaults to all frames.
Optional planner input name (start, goal, target, path_goal, or path_target).
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Prepare kinematics and geometry required by the task.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Accumulate normal equations directly from existing sparse rows.
Set the valid trajectory length for each batch element.
Per-batch valid frame counts.
Build for the trajectory in var_values; a built per-pose task keeps its identity and live targets.
A terminal cost whose dense task writes the trajectory layout itself (fill_terminal) takes one launch.
Match world-frame link targets; "distance" scores translation/angle thresholds, while "log" uses log residuals.
Example
>>> from robot_descriptions.loaders.yourdfpy import load_robot_description
>>> from robokit.opt.var_values import VarKey, VarValues
>>> from robokit.robo.robot import Robot
>>> robot = Robot.load(load_robot_description("ur10_description"))
>>> target_link_index = robot.spec.link_names.index("ee_link")
>>> state = robot.state(q=robot.spec.zero_q)
>>> target_link_pose = robot.forward_kinematics(state).get_T_world_link(target_link_index)
>>> frame_task = FrameTask(robot, target_link_index, target_link_pose, position_weight=1.0, orientation_weight=1.0)
>>> random_q = wp.from_numpy(np.array([[-0.5, 0.3, -0.2, 0.4, -0.1, 0.25]], dtype=np.float32), dtype=wp.float32, device="cpu")
>>> state = robot.state(q=random_q)
>>> robot.forward_kinematics(state)
>>> robot.compute_motion_subspace(state)
>>> weighted_residual = frame_task.compute_weighted_residual(VarValues({VarKey.ROBOT: state}))
>>> weighted_jacobian = frame_task.compute_weighted_jacobian(VarValues({VarKey.ROBOT: state}))
>>> assert np.allclose(weighted_residual.numpy()[0, :3], np.array([0.5839, -0.0795, -0.1784], dtype=np.float32), atol=1e-3)
>>> assert np.allclose(weighted_jacobian.numpy()[0, :, 0], np.array([-1.1628, 0.1626, -0.167, -0.0765, 0.3698, 0.9521], dtype=np.float32), atol=1e-3)import numpy as np
import warp as wp
from robokit.terms.dense.frame_task import FrameTask
from robot_descriptions.loaders.yourdfpy import load_robot_description
from robokit.opt.var_values import VarKey, VarValues
from robokit.robo.robot import Robot
robot = Robot.load(load_robot_description("ur10_description"))
target_link_index = robot.spec.link_names.index("ee_link")
state = robot.state(q=robot.spec.zero_q)
target_link_pose = robot.forward_kinematics(state).get_T_world_link(target_link_index)
frame_task = FrameTask(robot, target_link_index, target_link_pose, position_weight=1.0, orientation_weight=1.0)
random_q = wp.from_numpy(np.array([[-0.5, 0.3, -0.2, 0.4, -0.1, 0.25]], dtype=np.float32), dtype=wp.float32, device="cpu")
state = robot.state(q=random_q)
robot.forward_kinematics(state)
robot.compute_motion_subspace(state)
weighted_residual = frame_task.compute_weighted_residual(VarValues({VarKey.ROBOT: state}))
weighted_jacobian = frame_task.compute_weighted_jacobian(VarValues({VarKey.ROBOT: state}))
assert np.allclose(weighted_residual.numpy()[0, :3], np.array([0.5839, -0.0795, -0.1784], dtype=np.float32), atol=1e-3)
assert np.allclose(weighted_jacobian.numpy()[0, :, 0], np.array([-1.1628, 0.1626, -0.167, -0.0765, 0.3698, 0.9521], dtype=np.float32), atol=1e-3)False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare kinematics and geometry required by the task.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Allocate Warp buffers on device; a task without any keeps none.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Update target poses in-place. Batch size must match the original.
Complete a stub with the helper's links and target buffer (not copied); named weights pick those links.
Resolve scalar, ordered, or named template weights for the selected frames.
>>> task = FrameTask(position_weight={"tool": 2.0}, orientation_weight=3.0)
>>> task.weights_for(["tool"])[0].tolist()
[2.0]from robokit.terms.dense.frame_task import FrameTask
task = FrameTask(position_weight={"tool": 2.0}, orientation_weight=3.0)
task.weights_for(["tool"])[0].tolist()Give a config-time stub its robot, link indices, and targets, keeping the configured weights.
Write terminal pose residuals and optional sparse Jacobian values in one launch.
Layouts:
var_values: flattened poses,trajectory_framesper trajectory.- Jacobian: TrajectoryTask's single-frame, row-major values, one per actuated joint.
Match the world positions of one or more link frames.
Example
>>> from robot_descriptions.loaders.yourdfpy import load_robot_description
>>> from robokit.opt.var_values import VarKey, VarValues
>>> from robokit.robo.robot import Robot
>>> robot = Robot.load(load_robot_description("ur10_description"))
>>> target_link_index = robot.link_names.index("ee_link")
>>> state = robot.state(q=robot.zero_q)
>>> target_link_pose = robot.forward_kinematics(state).get_T_world_link(target_link_index).reshape((1, 1))
>>> position_task = PositionTask(robot, target_link_index, target_link_pose, weight=1.0)
>>> random_q = wp.from_numpy(np.array([[-0.5, 0.3, -0.2, 0.4, -0.1, 0.25]], dtype=np.float32), dtype=wp.float32)
>>> state = robot.state(q=random_q)
>>> robot.forward_kinematics(state)
>>> robot.compute_motion_subspace(state)
>>> weighted_residual = position_task.compute_weighted_residual(VarValues({VarKey.ROBOT: state}))
>>> weighted_jacobian = position_task.compute_weighted_jacobian(VarValues({VarKey.ROBOT: state}))
>>> weighted_residual.numpy().shape[1] == 3
True
>>> weighted_jacobian.numpy().shape[1] == 3
Trueimport numpy as np
import warp as wp
from robokit.terms.dense.position_task import PositionTask
from robot_descriptions.loaders.yourdfpy import load_robot_description
from robokit.opt.var_values import VarKey, VarValues
from robokit.robo.robot import Robot
robot = Robot.load(load_robot_description("ur10_description"))
target_link_index = robot.link_names.index("ee_link")
state = robot.state(q=robot.zero_q)
target_link_pose = robot.forward_kinematics(state).get_T_world_link(target_link_index).reshape((1, 1))
position_task = PositionTask(robot, target_link_index, target_link_pose, weight=1.0)
random_q = wp.from_numpy(np.array([[-0.5, 0.3, -0.2, 0.4, -0.1, 0.25]], dtype=np.float32), dtype=wp.float32)
state = robot.state(q=random_q)
robot.forward_kinematics(state)
robot.compute_motion_subspace(state)
weighted_residual = position_task.compute_weighted_residual(VarValues({VarKey.ROBOT: state}))
weighted_jacobian = position_task.compute_weighted_jacobian(VarValues({VarKey.ROBOT: state}))
weighted_residual.numpy().shape[1] == 3
weighted_jacobian.numpy().shape[1] == 3False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare kinematics and geometry required by the task.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Allocate Warp buffers on device; a task without any keeps none.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Update residual weights in-place. Accepts scalar or per-frame sequence.
Give a config-time stub its robot, link indices, and targets, keeping the configured weights.
Update target poses in-place. Batch size must match the original.
Complete a stub with the helper's links and target buffer, which it reads directly.
Match the world rotations of one or more link frames.
Example
>>> from robot_descriptions.loaders.yourdfpy import load_robot_description
>>> from robokit.opt.var_values import VarKey, VarValues
>>> from robokit.robo.robot import Robot
>>> robot = Robot.load(load_robot_description("ur10_description"))
>>> target_link_index = robot.link_names.index("ee_link")
>>> state = robot.state(q=robot.zero_q)
>>> target_link_pose = robot.forward_kinematics(state).get_T_world_link(target_link_index).reshape((1, 1))
>>> rotation_task = RotationTask(robot, target_link_index, target_link_pose, weight=1.0)
>>> random_q = wp.from_numpy(np.array([[-0.5, 0.3, -0.2, 0.4, -0.1, 0.25]], dtype=np.float32), dtype=wp.float32)
>>> state = robot.state(q=random_q)
>>> robot.forward_kinematics(state)
>>> robot.compute_motion_subspace(state)
>>> weighted_residual = rotation_task.compute_weighted_residual(VarValues({VarKey.ROBOT: state}))
>>> weighted_jacobian = rotation_task.compute_weighted_jacobian(VarValues({VarKey.ROBOT: state}))
>>> weighted_residual.numpy().shape[1] == 3
True
>>> weighted_jacobian.numpy().shape[1] == 3
Trueimport numpy as np
import warp as wp
from robokit.terms.dense.rotation_task import RotationTask
from robot_descriptions.loaders.yourdfpy import load_robot_description
from robokit.opt.var_values import VarKey, VarValues
from robokit.robo.robot import Robot
robot = Robot.load(load_robot_description("ur10_description"))
target_link_index = robot.link_names.index("ee_link")
state = robot.state(q=robot.zero_q)
target_link_pose = robot.forward_kinematics(state).get_T_world_link(target_link_index).reshape((1, 1))
rotation_task = RotationTask(robot, target_link_index, target_link_pose, weight=1.0)
random_q = wp.from_numpy(np.array([[-0.5, 0.3, -0.2, 0.4, -0.1, 0.25]], dtype=np.float32), dtype=wp.float32)
state = robot.state(q=random_q)
robot.forward_kinematics(state)
robot.compute_motion_subspace(state)
weighted_residual = rotation_task.compute_weighted_residual(VarValues({VarKey.ROBOT: state}))
weighted_jacobian = rotation_task.compute_weighted_jacobian(VarValues({VarKey.ROBOT: state}))
weighted_residual.numpy().shape[1] == 3
weighted_jacobian.numpy().shape[1] == 3False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare kinematics and geometry required by the task.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Allocate Warp buffers on device; a task without any keeps none.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Update residual weights in-place. Accepts scalar or per-frame sequence.
Give a config-time stub its robot, link indices, and targets, keeping the configured weights.
Update target poses in-place. Batch size must match the original.
Complete a stub with the helper's links and target buffer (not copied); named weights zero unlisted links.
Drive point-pair deltas to target vectors; weight_mode "cost" takes sqrt, "normalized" divides by sqrt(rows).
Direction metrics:
- Default: robot-length-scaled vector error.
- `
direction_only=Truewithdirection_metric="cosine":1 - dot(normalized_robot, normalized_target)`.
Named weights pick their rows when the task is built:
- A target name tracks that target's row of
targets, ordered bytarget_names, from the world origin to its link intarget_links. (origin link, task link, origin target, task target)reads the pair roworigin * len(target_names) + taskof a pair buffer.
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare kinematics and geometry required by the task.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Attach the robot and helper links/targets, reading Warp buffers in place; named weights pick the rows.
Update targets (and optionally weights) in place; shapes must match the built buffers.
Penalize world-frame CoM target error using mass-weighted link CoM positions and their analytic Jacobians.
Example
>>> from robot_descriptions.loaders.yourdfpy import load_robot_description
>>> from robokit.robo.robot import Robot
>>> robot = Robot.load(load_robot_description("g1_description"))
>>> task = ComPositionTask(robot=robot, target_com_position=np.zeros(3), weight=1.0)
>>> task.residual_dim
3import numpy as np
from robokit.terms.dense.com_position_task import ComPositionTask
from robot_descriptions.loaders.yourdfpy import load_robot_description
from robokit.robo.robot import Robot
robot = Robot.load(load_robot_description("g1_description"))
task = ComPositionTask(robot=robot, target_com_position=np.zeros(3), weight=1.0)
task.residual_dimFalse when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Label a config gave this term, so config.task(name) finds it again.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Swap placeholders, then take the robot of the state this task reads.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare kinematics and geometry required by the task.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Allocate Warp buffers on device; a task without any keeps none.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Update CoM targets while retaining allocated buffers and captured graphs.
Example
>>> task = ComPositionTask(target_com_position=np.zeros(3)) >>> task.set_target(np.array([0.1, 0.0, 0.8])) >>> task.target_com_position.shape (3,)
import numpy as np from robokit.terms.dense.com_position_task import ComPositionTask task = ComPositionTask(target_com_position=np.zeros(3)) task.set_target(np.array([0.1, 0.0, 0.8])) task.target_com_position.shape
Update the three CoM residual weights in place.
Example
>>> task = ComPositionTask(target_com_position=np.zeros(3)) >>> task.set_weight([1.0, 1.0, 0.0]) >>> task.weight [1.0, 1.0, 0.0]
import numpy as np from robokit.terms.dense.com_position_task import ComPositionTask task = ComPositionTask(target_com_position=np.zeros(3)) task.set_weight([1.0, 1.0, 0.0]) task.weight
Penalize joint and optional floating-base translation limit violations.
Example
>>> from robot_descriptions.loaders.yourdfpy import load_robot_description
>>> from robokit.opt.var_values import VarKey, VarValues
>>> from robokit.robo.robot import Robot
>>> robot = Robot.load(load_robot_description("ur10_description"))
>>> joint_var = robot.spec.actuated_joint_limits[:, 0] - 0.1
>>> state = robot.state(q=wp.from_numpy(joint_var[None].astype(np.float32), dtype=wp.float32))
>>> position_limit = PositionLimit(robot=robot)
>>> weighted_residual = position_limit.compute_weighted_residual(VarValues({VarKey.ROBOT: state}))
>>> np.allclose(weighted_residual.numpy()[0], np.array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1], dtype=np.float32))
True
>>> weighted_jacobian = position_limit.compute_weighted_jacobian(VarValues({VarKey.ROBOT: state}))
>>> np.allclose(np.diag(weighted_jacobian.numpy()[0]), np.array([-1.0, -1.0, -1.0, -1.0, -1.0, -1.0], dtype=np.float32))
Trueimport numpy as np
import warp as wp
from robokit.terms.dense.position_limit import PositionLimit
from robot_descriptions.loaders.yourdfpy import load_robot_description
from robokit.opt.var_values import VarKey, VarValues
from robokit.robo.robot import Robot
robot = Robot.load(load_robot_description("ur10_description"))
joint_var = robot.spec.actuated_joint_limits[:, 0] - 0.1
state = robot.state(q=wp.from_numpy(joint_var[None].astype(np.float32), dtype=wp.float32))
position_limit = PositionLimit(robot=robot)
weighted_residual = position_limit.compute_weighted_residual(VarValues({VarKey.ROBOT: state}))
np.allclose(weighted_residual.numpy()[0], np.array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1], dtype=np.float32))
weighted_jacobian = position_limit.compute_weighted_jacobian(VarValues({VarKey.ROBOT: state}))
np.allclose(np.diag(weighted_jacobian.numpy()[0]), np.array([-1.0, -1.0, -1.0, -1.0, -1.0, -1.0], dtype=np.float32))False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare kinematics and geometry required by the task.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Build from the robot; a base bound needs a floating base, e.g. IK with enable_T_world_base=True.
Update residual weights in-place. Accepts a scalar or per-joint sequence.
Limit the angle between a frame-local axis and a world-frame axis.
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Label a config gave this term, so config.task(name) finds it again.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Swap placeholders, then take the robot of the state this task reads.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare kinematics and geometry required by the task.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Device of the task's Warp buffers; None before init_buffers.
Allocate Warp buffers on device; a task without any keeps none.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Penalize joint velocity limit violations relative to a previous state.
Example
>>> from robot_descriptions.loaders.yourdfpy import load_robot_description
>>> from robokit.robo.robot import Robot
>>> robot = Robot.load(load_robot_description("panda_description"))
>>> state_prev = robot.state(q=robot.spec.zero_q)
>>> state_curr = robot.state(q=robot.spec.zero_q + 0.5)
>>> vel_task = VelocityLimitTask(robot=robot, dt=0.1, weight=1.0)
>>> # Residual should be positive if velocity exceeds limitsfrom robokit.terms.dense.velocity_limit_task import VelocityLimitTask
from robot_descriptions.loaders.yourdfpy import load_robot_description
from robokit.robo.robot import Robot
robot = Robot.load(load_robot_description("panda_description"))
state_prev = robot.state(q=robot.spec.zero_q)
state_curr = robot.state(q=robot.spec.zero_q + 0.5)
vel_task = VelocityLimitTask(robot=robot, dt=0.1, weight=1.0)
# Residual should be positive if velocity exceeds limitsFalse when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare shared state before this term is evaluated.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Build from the robot; a previous state from the helper is its buffer, read directly.
Update residual weights in-place. Accepts a scalar or per-joint sequence.
Penalize robot collision spheres that approach scene geometry.
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Swap placeholders, then take the robot of the state this task reads.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare kinematics and geometry required by the task.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Device of the task's Warp buffers; None before init_buffers.
Allocate Warp buffers on device; a task without any keeps none.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Whether the term adds cost: nonzero weight and a scene with geometry, or a helper's scene to come.
Build a config-time stub, with fresh query buffers.
Copy a per-query device scene mapping into stable query buffers.
Compute collision cost and optional direct gradient.
Lifecycle
- Prepare trajectory-shaped collision buffers and kinematics.
- Query the scene SDF for every selected sphere.
- Accumulate the selected cost form and optional gradient.
Minimize hand contact-point distance to a scene surface.
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Label a config gave this term, so config.task(name) finds it again.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Swap placeholders, then take the robot of the state this task reads.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare kinematics and geometry required by the task.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Device of the task's Warp buffers; None before init_buffers.
Allocate Warp buffers on device; a task without any keeps none.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Compute cost and optional gradient.
Lifecycle
- Transform local contact points or reuse supplied world points.
- Query or reuse closest points, signed distances, and normals.
- Compute residuals and accumulate cost and
Jᵀr.
Penalize the most active robot self-collision pairs.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Swap placeholders, then take the robot of the state this task reads.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare kinematics and geometry required by the task.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Device of the task's Warp buffers; None before init_buffers.
Allocate Warp buffers on device; a task without any keeps none.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Build a config-time stub's collision pairs with the configured settings.
Compute self-collision residuals.
Lifecycle
- Prepare world collision geometry.
- Select the most active broad-phase link pairs.
- Evaluate the closest sphere or capsule pair for each slot.
Compute the analytic self-collision Jacobian.
Lifecycle
- Prepare kinematics and motion subspaces.
- Reuse the residual's active-pair selection when available.
- Differentiate the selected sphere or capsule distances.
Compute fused link-sphere cost and optional gradient.
Lifecycle
- Prepare link spheres and motion subspaces.
- Allocate a scratch gradient for cost-only evaluations.
- Accumulate every active link-pair contribution.
Penalize change from a previous joint and optional base state.
Example
>>> from robot_descriptions.loaders.yourdfpy import load_robot_description
>>> from robokit.opt.var_values import VarKey
>>> from robokit.robo.robot import Robot
>>> robot = Robot.load(load_robot_description("panda_description"))
>>> state_prev = robot.state(q=robot.spec.zero_q)
>>> state_curr = robot.state(q=robot.spec.zero_q + 0.1)
>>> smooth_task = SmoothnessTask(robot=robot, weight=0.5)
>>> values = VarValues({VarKey.ROBOT: state_curr})
>>> cost = wp.zeros(1, device=state_curr.q.device)
>>> gradient = wp.zeros((1, state_curr.tangent_dim), device=state_curr.q.device)
>>> smooth_task.compute_weighted_cost_and_gradient(values, state_prev, out_cost=cost, out_gradient=gradient)
>>> np.allclose(gradient.numpy(), 0.025)
Trueimport numpy as np
import warp as wp
from robokit.terms.dense.smoothness_task import SmoothnessTask
from robokit.opt.var_values import VarValues
from robot_descriptions.loaders.yourdfpy import load_robot_description
from robokit.opt.var_values import VarKey
from robokit.robo.robot import Robot
robot = Robot.load(load_robot_description("panda_description"))
state_prev = robot.state(q=robot.spec.zero_q)
state_curr = robot.state(q=robot.spec.zero_q + 0.1)
smooth_task = SmoothnessTask(robot=robot, weight=0.5)
values = VarValues({VarKey.ROBOT: state_curr})
cost = wp.zeros(1, device=state_curr.q.device)
gradient = wp.zeros((1, state_curr.tangent_dim), device=state_curr.q.device)
smooth_task.compute_weighted_cost_and_gradient(values, state_prev, out_cost=cost, out_gradient=gradient)
np.allclose(gradient.numpy(), 0.025)False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare shared state before this term is evaluated.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Build from the robot; a previous state from the helper is its buffer, read directly.
Update joint-block residual weights in-place. Base-block weights are preserved.
Accumulate half the squared residual norm and its direct gradient.
Bias joints/base toward rest; `rest_q="midrange"` uses joint-limit midpoints as the rest state.
Supported variables:
- RobotState: joints and optional floating base.
- VectorVar: joint-only block, sized by
rest_qwith no robot. - SE3Var: pose-only block, including model-offset regularization.
Example
>>> from robot_descriptions.loaders.yourdfpy import load_robot_description
>>> from robokit.opt.var_values import VarKey, VarValues
>>> from robokit.robo.robot import Robot
>>> robot = Robot.load(load_robot_description("panda_description"))
>>> rest_q = robot.spec.midrange_q
>>> configuration_task = ConfigurationTask(robot=robot, rest_q=rest_q, weight=0.1)
>>> q_init = wp.from_numpy(rest_q + 0.1, dtype=wp.float32)
>>> state = robot.state(q=q_init)
>>> residual = configuration_task.compute_weighted_residual(VarValues({VarKey.ROBOT: state}))
>>> # Residual should be ~0.1 * weightimport warp as wp
from robokit.terms.dense.configuration_task import ConfigurationTask
from robot_descriptions.loaders.yourdfpy import load_robot_description
from robokit.opt.var_values import VarKey, VarValues
from robokit.robo.robot import Robot
robot = Robot.load(load_robot_description("panda_description"))
rest_q = robot.spec.midrange_q
configuration_task = ConfigurationTask(robot=robot, rest_q=rest_q, weight=0.1)
q_init = wp.from_numpy(rest_q + 0.1, dtype=wp.float32)
state = robot.state(q=q_init)
residual = configuration_task.compute_weighted_residual(VarValues({VarKey.ROBOT: state}))
# Residual should be ~0.1 * weightFalse when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare shared state before this term is evaluated.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Build from the robot; Warp rest values from the helper are its buffers, read directly.
Update joint-block residual weights in-place. Base-block weights are preserved.
Update the rest references; a config-time stub keeps NumPy references until its buffers are built.
Maximize Yoshikawa sqrt(det(J @ J.T)) at frame_index, using the end-effector translational Jacobian.
Analytic derivative columns:
- Fixed-base joints: computed normally.
- Floating base: left zero.
Example
>>> import numpy as np
>>> from robot_descriptions.loaders.yourdfpy import load_robot_description
>>> from robokit.opt.var_values import VarKey, VarValues
>>> from robokit.robo.robot import Robot
>>> robot = Robot.load(load_robot_description("ur10_description"))
>>> frame_index = robot.spec.link_names.index("ee_link")
>>> random_q = wp.from_numpy(
... np.array([[-0.5, 0.3, -0.2, 0.4, -0.1, 0.25]], dtype=np.float32), dtype=wp.float32
... )
>>> state = robot.state(q=random_q)
>>> task = ManipulabilityTask(robot, frame_index, weight=0.01)
>>> task.compute_weighted_residual(VarValues({VarKey.ROBOT: state})).numpy().shape[1] == 1
True
>>> task.compute_weighted_jacobian(VarValues({VarKey.ROBOT: state})).numpy().shape[1] == 1
Trueimport warp as wp
from robokit.terms.dense.manipulability_task import ManipulabilityTask
import numpy as np
from robot_descriptions.loaders.yourdfpy import load_robot_description
from robokit.opt.var_values import VarKey, VarValues
from robokit.robo.robot import Robot
robot = Robot.load(load_robot_description("ur10_description"))
frame_index = robot.spec.link_names.index("ee_link")
random_q = wp.from_numpy(
np.array([[-0.5, 0.3, -0.2, 0.4, -0.1, 0.25]], dtype=np.float32), dtype=wp.float32
)
state = robot.state(q=random_q)
task = ManipulabilityTask(robot, frame_index, weight=0.01)
task.compute_weighted_residual(VarValues({VarKey.ROBOT: state})).numpy().shape[1] == 1
task.compute_weighted_jacobian(VarValues({VarKey.ROBOT: state})).numpy().shape[1] == 1False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Label a config gave this term, so config.task(name) finds it again.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Swap placeholders, then take the robot of the state this task reads.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare kinematics and geometry required by the task.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Allocate Warp buffers on device; a task without any keeps none.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Update the manipulability weight in-place (live-tunable, CUDA-graph safe).
Force closure task: contacts from var.contact_point_centers_world, normals from the SDF of warp_meshes.
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Label a config gave this term, so config.task(name) finds it again.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Swap placeholders, then take the robot of the state this task reads.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare kinematics and geometry required by the task.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Device of the task's Warp buffers; None before init_buffers.
Allocate Warp buffers on device; a task without any keeps none.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Read the contact points from this variable instead of the stored array.
A SceneDistanceTask on the same contacts, listed before this term: its scene query serves both.
Preserve interaction-mesh Laplacian coordinates from human to robot.
Example
>>> from robot_descriptions.loaders.yourdfpy import load_robot_description
>>> import numpy as np
>>> from robokit.helpers.humanoid_retargeting import build_interaction_mesh_frame
>>> from robokit.robo.robot import Robot
>>> robot = Robot.load(load_robot_description("panda_description"))
>>> state = robot.state(q=robot.zero_q)
>>> robot.forward_kinematics(state)
>>> links = [robot.link_names.index(n) for n in robot.link_names[1:5]]
>>> P = state.T_world_link.numpy()[0][links, :3]
>>> pts = np.vstack([P, P.mean(0) + np.array([[0.2, 0, 0], [0, 0.2, 0], [0, 0, 0.2]], np.float32)])
>>> task = InteractionMeshTask(robot, np.array(links, dtype=np.int32))
>>> T_world_object = np.array([0, 0, 0, 1, 0, 0, 0], np.float32)
>>> task.set_frame(*build_interaction_mesh_frame(P, pts[len(P):], T_world_object))
>>> task.residual_dim
21from robokit.terms.dense.interaction_mesh_task import InteractionMeshTask
from robot_descriptions.loaders.yourdfpy import load_robot_description
import numpy as np
from robokit.helpers.humanoid_retargeting import build_interaction_mesh_frame
from robokit.robo.robot import Robot
robot = Robot.load(load_robot_description("panda_description"))
state = robot.state(q=robot.zero_q)
robot.forward_kinematics(state)
links = [robot.link_names.index(n) for n in robot.link_names[1:5]]
P = state.T_world_link.numpy()[0][links, :3]
pts = np.vstack([P, P.mean(0) + np.array([[0.2, 0, 0], [0, 0.2, 0], [0, 0, 0.2]], np.float32)])
task = InteractionMeshTask(robot, np.array(links, dtype=np.int32))
T_world_object = np.array([0, 0, 0, 1, 0, 0, 0], np.float32)
task.set_frame(*build_interaction_mesh_frame(P, pts[len(P):], T_world_object))
task.residual_dimRobot model. May be assigned later by the solver.
Robot link for each dynamic vertex.
Residual multiplier.
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare kinematics and geometry required by the task.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Build from the robot and the helper's links, the first vertices of every frame.
Set adjacency, reference vertices, and target Laplacian for one frame.
Concatenate weighted residuals, with optional threshold-based selection scoring.
Modes:
- Without
thresholds: ordinary residual concatenation for optimization. - With
thresholds: bounded children are constraints; children with None thresholds are scores.
Selection sequence:
- Require each constraint's absolute residuals to satisfy scalar, per-row or repeating bounds.
- Return score norm for valid rows and infinity otherwise, preserving ordinary squared costs on valid rows.
- Count valid rows per
num_candidatesgroup invalid_count; use stable selection to preserve ties.
Example
>>> from robokit.terms.dense.configuration_task import ConfigurationTask >>> task = CompositeTask([ConfigurationTask()], thresholds=[None]) >>> task.residual_dim 1
from robokit.terms.composite_task import CompositeTask from robokit.terms.dense.configuration_task import ConfigurationTask task = CompositeTask([ConfigurationTask()], thresholds=[None]) task.residual_dim
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Label a config gave this term, so config.task(name) finds it again.
The robot a built task reads; None on a template.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Device of the task's Warp buffers; None before init_buffers.
Allocate Warp buffers on device; a task without any keeps none.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Build a copy of each unbuilt child, dropping one the context empties, and take the candidate count.
Prepare every child task.
Align robot-link positions to optionally indexed world-frame points.
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare kinematics and geometry required by the task.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Accumulate normal equations directly from existing sparse rows.
Write residuals and optional sparse Jacobian values into shared buffers.
Optimization variables.
Shared residual buffer.
First residual row for this task.
Optional shared sparse-value buffer.
First sparse entry for this task.
Set the valid trajectory length for each batch element.
Per-batch valid frame counts.
Fill a template with helper links and targets, reading a Warp target directly; unlisted links weigh zero.
Keep selected link points within a deadband of trajectory contacts: points, or segments along contact_axes.
Robot with FK support.
Shared [T, C, 3] or batched [B, T, C, 3] world targets.
Padded [T, C] robot link indices.
Shared [T, C] or batched [B, T, C] validity mask.
Deadband margin; no penalty when distance is within this threshold.
Residual weight.
Optional link-local points. Defaults to each link origin.
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare kinematics and geometry required by the task.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Allocate Warp buffers on device; a task without any keeps none.
Accumulate normal equations directly from existing sparse rows.
Write residuals and optional sparse Jacobian values into shared buffers.
Optimization variables.
Shared residual buffer.
First residual row for this task.
Optional shared sparse-value buffer.
First sparse entry for this task.
Set the valid trajectory length for each batch element.
Per-batch valid frame counts.
Build from the helper's contact links; its Warp contact buffers are read directly.
Copy batched runtime contact data into graph-stable device buffers; omitted points and axes are zero. A nonzero axis makes its contact the segment from the point along it: the link point may sit anywhere on it, e.g. a straight move at its own pace.
Evaluate trajectory-sphere costs against geometry, else every scene geometry; one row block per geometry.
Swept-segment evaluation:
- Linearly interpolate sphere centers at
sweep_stepssamples between waypoints. - Differentiate the largest residual through both endpoints.
- Include the final segment even if
free_goal_framedisables the standalone goal waypoint.
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare kinematics and geometry required by the task.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Accumulate normal equations directly from existing sparse rows.
Set the valid trajectory length for each batch element.
Per-batch valid frame counts.
Whether the term adds cost: nonzero weight, a scene, and at least one selected sphere.
Build for the robot and horizon in var_values, sharing a Warp frame_sphere_mask, unless inactive.
Build stable scene/sparsity parameters; query caches grow on first use.
Lifecycle
- Upload selected sphere and actuated-column indices.
- Build the actuated-column joint CSR mapping.
- Allocate scene/margin arrays and invalidate lazy query caches.
Max-accumulate penetration from the most recent Jacobian query cache, over every geometry.
Runs over the batch the cache holds: out may be longer (a score task sized for every draw).
Set the active (batch, frame, selected sphere) collision rows.
One block of rows per geometry; an empty scene keeps one block of zeros.
Set scenes shaped (batch,) or per-frame scenes shaped (batch, frames).
Compute sparse cost and optional gradient with one scene query.
Lifecycle
- Use the base cost-only path when no gradient is requested.
- Compute Jacobian values and cache their scene query.
- Rebuild residuals from the cache and reduce cost and gradient.
Fill residuals and sparse values from one shared scene query.
Compute trajectory collision residuals.
Lifecycle
- Prepare forward kinematics and stable runtime buffers.
- Run each geometry's primitive, mesh, volume, or hybrid query kernel.
- Evaluate every frame and selected sphere.
Compute sparse trajectory collision Jacobian values.
Lifecycle
- Prepare kinematics, motion subspaces, and runtime-sized caches.
- Query each geometry and cache its sphere centers and SDF normals.
- Assemble active joint and optional floating-base columns.
Penalize self-collisions between robot collision spheres at all frames.
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare kinematics and geometry required by the task.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Accumulate normal equations directly from existing sparse rows.
Write residuals and optional sparse Jacobian values into shared buffers.
Optimization variables.
Shared residual buffer.
First residual row for this task.
Optional shared sparse-value buffer.
First sparse entry for this task.
Set the valid trajectory length for each batch element.
Per-batch valid frame counts.
Build from the robot and horizon of the trajectory var_values holds.
Penalize joint differences and optional floating-base motion.
Configuration options:
- robot/num_frames omitted:
resolvetakes them from the trajectory it reads. - smooth_boundary=False: disable the first and last finite-difference rows.
- history: an executed
(q_history, history_valid)prefix, e.g. a helper's shared buffers.
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare shared state before this term is evaluated.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Accumulate normal equations directly from existing sparse rows.
Write residuals and optional sparse Jacobian values into shared buffers.
Optimization variables.
Shared residual buffer.
First residual row for this task.
Optional shared sparse-value buffer.
First sparse entry for this task.
Set the valid trajectory length for each batch element.
Per-batch valid frame counts.
Return [K,T,J] paths, preserving alternative initial pose hypotheses; each initial hypothesis retains a distinct branch, and the output contains at most count paths; periodic joints use angular distances and return continuous equivalent angles within supplied limits; if no continuous lift fits the limits, the original bounded path is retained.
>>> q = np.array([[[0.], [3.]], [[3.], [0.1]], [[0.2], [3.]]]) >>> TrajectorySmoothnessTask().compute_paths(q, np.array([[0., .1], [.1, 0.], [0., .1]]), np.arange(3.)).ravel().tolist() [0.0, 0.1, 0.2]
import numpy as np from robokit.terms.sparse.trajectory_smoothness_task import TrajectorySmoothnessTask q = np.array([[[0.], [3.]], [[3.], [0.1]], [[0.2], [3.]]]) TrajectorySmoothnessTask().compute_paths(q, np.array([[0., .1], [.1, 0.], [0., .1]]), np.arange(3.)).ravel().tolist()
Sample constant joint shifts that preserve every frame's joint limits; full-turn joints need at most one revolution, and extra cubic samples retain small shifts.
Build finite-difference weights and history buffers.
Lifecycle
- Build joint and optional base weights.
- Build the requested difference stencil.
- Allocate history and reference trajectories.
Build from the robot and horizon of the trajectory var_values holds, and the helper's values.
Set the fixed executed prefix. No-op at order=1, which never reads it.
Penalize bound violations with `weight * max(lower - derivative, derivative - upper, 0)`.
Derivative orders:
- 0: position.
- 1: velocity.
- 2: acceleration.
- 3: jerk.
Difference schemes:
- `
central`: retimer's five-point derivatives with held endpoints; multiple orders share one task and launch. - `
forward`: consecutive finite differences without padded endpoints; one order only.
Broadcast layouts:
- Bounds/weights: broadcast over joints, with a leading order axis when `
order` is a sequence. - Scalar dt: shared by all plans.
- Warp-vector dt: one per plan, repeated in instance-major order for candidate batches.
Time scaling:
- Default: divide each derivative by dt raised to its order.
- Explicit `
dt_power`: override the exponent per order. - `
timing = (frame_segment, segments, time_weight): a bound row's step is its plan'sexp(VarKey.SEGMENT_STEP)` for its frame's segment, a Jacobian column (zero-width bounds keep `dt);time_weight` > 0 costs each second.
Unbuilt templates (resolve):
- `
upper` may hold one entry per order; unset bounds default to the robot's position and velocity limits. - `
lowerdefaults to the robot's lower position limit and to-upper` for derivatives. - `
activation` keeps that fraction of each range away from both bounds, indexed by order. - `
normalize` divides the weights by the bounds.
Example
>>> task = TrajectoryLimitTask(robot, 33, order=3, lower=-500.0, upper=500.0) # doctest: +SKIP
from robokit.terms.sparse.trajectory_limit_task import TrajectoryLimitTask task = TrajectoryLimitTask(robot, 33, order=3, lower=-500.0, upper=500.0) # doctest: +SKIP
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare shared state before this term is evaluated.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Accumulate normal equations directly from existing sparse rows.
Set the valid trajectory length for each batch element.
Per-batch valid frame counts.
Build a template's robot, horizon, step and bounds from the trajectory var_values holds.
Fuse central bounds at dt and all-frame abs position limits into one launch, skipping tasks in keep.
Set a positive scalar time step or a float32 Warp vector with one entry per plan.
Write the residual and, when asked, the Jacobian values in one launch.
Pairwise position and direction residuals for trajectory retargeting.
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare kinematics and geometry required by the task.
Add each row's weighted cost into out_cost, for a term without residuals (a simulated rollout, or one of its per-step costs); PopulationSolver calls it on every term that is not a ResidualTask.
Accumulate normal equations directly from existing sparse rows.
Write residuals and optional sparse Jacobian values into shared buffers.
Optimization variables.
Shared residual buffer.
First residual row for this task.
Optional shared sparse-value buffer.
First sparse entry for this task.
Set the valid trajectory length for each batch element.
Per-batch valid frame counts.
Build from the helper's points, links, pairs and scales; Warp points and scales are read directly.
Build target, pair, and sparse ancestor buffers.
Lifecycle
- Upload targets and pair indices.
- Build selected-link ancestor masks.
- Collect the active actuated columns.