Get the dimensions of the action space. Returns: action_space_dims: the action space dimensions

Construct the dense matrix D^T s D for multiple orders of smoothing. Explanation of the smoothing lambda term: For 1st/2nd/3rd smoothing ter

Convert a token to a special token.

Transform the action space to the trajectory. Args: action: (..., *action_space_dims) traj_history_xyz: (..., T, 3)

Estimate the t0 states from the trajectory history.

Computes the yaw angle given an so3 rotation matrix (assumes that rotation is described in xyz order) Args: rot_mat (torch.Tensor): [

Solve the following problem: min_x={x_0, ..., x_N-1} sum_i={0, ..., N-1} w_data_i (x_i * s_i - y_i)**2 + smooth_terms Args: s (...,

Unwrap the last dimension of the tensor to make sure the diff is in (-pi, pi].

Build the processor with trajectory tokens.

Build tokenizer with trajectory tokens.

Euler integration for flow matching. Args: batch_size: The batch size. step_fn: The denoising step function.

Initialize Qwen3-VL VLM backbone. Qwen3-VL uses Qwen3VLForConditionalGeneration from transformers. See: https://huggingface.co/Qwen/Q

Initialize trajectory tokenizers.

Initialize the VLM backbone based on configuration.

Initialize VLM-specific configuration based on backend type.

Normalize the input tensor.

Recursively set attribute on object and all its children.

Compute the curvature from the theta, v, a, jerk. The kappa is computed by s = dt * v + dt^2 * a / 2 kappa = dtheta

Compute the acceleration from the velocity. Define: Δv_t = v_t+1 - v_t According to the kinematic model Δv_t

Validate that all required mixin attributes are present.

Given the dxy and theta, compute the velocity. The velocity is defined by the trapezoidal integration: define: u_t = [cos theta_t, si

Given the dxy and theta, compute the velocity. The velocity is defined by the trapezoidal integration: define: u_t = [cos theta_t, si

Encodes the trajectories as discrete tokens. The model conditions on the historical waypoints to tokenize the future waypoints. Trajectories c

Extract text between special tokens. Args: decoded_batch (list[str]): the data we extract from token (str): the special token we

Extract text tokens (e.g., cot, meta_action, question, answers) from the output tokens. Args: output_tokens (torch.Tensor): The output to

Build the banded matrix for the first-order smoothing term.

Fuse the trajectory tokens into the input ids. Args: input_ids: [B, n_token] traj_data: dict containing ego_history_x

Get the rank (GPU device) of the worker. Returns: rank (int): The rank of the worker.

Calculate yaw rotation matrices using polynomial fitting for both x(t) and y(t) Args: trajectory: np.array of shape (B, N, 3) for batch o

Load data from physical_ai_av for model inference. This function loads a sample from the physical_ai_av dataset and converts it to the format

Add a prefix with the rank to a message.

Replace pad tokens in input_ids with new token values.

Overwrite token IDs after the first EOS token with padding token ID. Args: token_ids (torch.Tensor): Token IDs of shape [B, L]. e

Converts a 2D rotation matrix to a 3D rotation matrix assuming flat xy plane. Args: rot (torch.Tensor): The 2D rotation matrix to convert

Creates one or many 2D rotation matrices. Args: angle (Union[float, np.ndarray]): The angle to rotate points by. if float, re

Creates one or many 2D rotation matrices. Args: angle (torch.Tensor): The angle to rotate points by. Size: [...]. Returns: t

Normalize angles to the range [-pi, pi] in torch. Args: x: Angle(s) in radians, can be numpy array or torch tensor Returns:

Sample from the diffusion model. Args: batch_size: The batch size. step_fn: The denoising step function that takes a

Sample data from the model. Args: batch_size: The batch size. step_fn: The denoising step function. devic

Sample trajectories from the data with VLM rollout. Args: data: The input data. top_p: The top-p value for sampling.

Build the banded matrix for the second-order smoothing term.

Solve a single-point constrained sequence with multiple orders of smoothing. This function solves the following problem: min_x={x_1, ...,

Smooth the heading of the trajectory. Args: traj_future_rot: (..., T, 3, 3)

Build the banded matrix for the third-order smoothing term.

Tokenize the history trajectory with prefix shape of (B, n_traj, ...). Args: tokenizer: Trajectory tokenizer with encode method t

Transform the future trajectory to the action space. Args: traj_history_xyz: (..., T, 3) traj_history_rot: (..., T, 3

Denoising step function. Args: x: The input tensor. t: The timestep. Returns: torch.Tensor: The

Call the stopping criteria. Args: input_ids (torch.LongTensor): The input IDs of shape [B, L]. scores (torch.FloatTen

Call the ExpertLogitsProcessor to mask out the logits for discrete trajectory tokens. The discrete trajectory tokens are not used for the exp

Initialize the BaseDiffusion model. Args: x_dims: The dimension of the input tensor.

Initialize the FlowMatching model. Args: int_method: The integration method used in inference. num_inference_steps: T

Initialize the UnicycleAccelCurvatureActionSpace. Args: accel_mean: Mean for normalizing acceleration. accel_std: Std

Initializes the tokenizer.

Initializes a multi-GPU-friendly python command line logger that logs on all processes with their rank prefixed in the log message. A

Initializes the tokenizer.

Args: eos_token_id (int): The EOS token ID.

Initialize the ExpertLogitsProcessor. Args: traj_token_offset: The offset of the trajectory tokens. traj_vocab_size:

Initialize the Fourier encoder V2. Args: dim: Output dimension of the encoder. Must be even as it's split into si

Initialize the per-waypoint action projection module V2. Args: in_dims: List of input dimensions. The last element specifies the

Transform the action space to the trajectory. Args: action: (..., T, 2) traj_history_xyz: (..., T, 3) tra

This function wraps angles to lie within [-pi, pi). Args: radians (np.ndarray): The input array of angles (in radians). Returns:

Training step for the flow matching model.

Construct the training data for the flow matching model.

Construct the message using images and cot.

Decodes the given tokens into future trajectories. We assume the the future tstamp is consistent with the future trajectory. Args:

Decodes the given tokens into future trajectories. The future trajectory is returned in the same coordinate frame as the historical trajectory

Encodes the trajectories as discrete tokens. We assume the the future tstamp is consistent with the future trajectory. Args:

Converts the euler angles representation to the so3 rotation matrix Args: euler_angles (np.array): euler angles [n,3] degrees bool

Extract the trajectory tokens from the output tokens (parallel/vectorized version). This is a fully vectorized implementation that processes all

Normalize the input tensor.

(B, C) -> (B, outdim)

Forward pass of the Fourier encoder V2. Args: x: Input tensor of arbitrary shape (..., ). Returns: Fourier-e

Forward pass of the per-waypoint action projection V2. Args: x: Action tensor of shape (batch_size, num_waypoints, action_dim).

Load submodules with pretrained submodules and initialize the model.

Get the dimensions of the action space.

Get the input embeddings of the model.

Get the output embeddings of the model.

Get the processor for the Qwen3-VL-2B-Instruct model.

Check if a normalized action is within bounds. Args: action: (..., N, 2) Returns: is_within_bounds: (...,)

Check if the action is within the bounds. By default, we assume the action is within bounds (dummy implementation). Args:

Delegate a log call to the underlying logger, after prefixing its message with the rank of the process it's being logged from. If `'rank'` is

Robust arctan2 for pytorch torch.arctan2(0,0)=nan, this function avoids the nan situation and returns ratan2(0,0)=0

Converts a 3D rotation matrix to a 2D rotation matrix by taking the x and y axes of the 3D rotation matrix, projecting them to xy plan, and perfor

Normalize angles to the range [-pi, pi]. Args: x: Angle(s) in radians, can be numpy array or torch tensor Returns: Normalize

Configure the root logger with a single colored console handler. Clears existing root handlers first to avoid duplicate log output when used

Computes the yaw angle given an so3 rotation matrix (assumes that rotation is described in xyz order) Args: rot_mat (np.ndarray): [..

Orthonormalize two 3D vectors using a stable Gram-Schmidt step. Args: M: Tensor of shape (..., 3, 2) with vectors (x, y). Returns:

Delegate weight tying to the nested VLM model.

Recursively cast data into the specified device, dtype.

Transform the future trajectory to the action space. Here we assume the traj_history_xyz[..., -1, :] is the current position and is all zeros

Args: coords (np.ndarray): [..., 2] coordinates offset (Optional[np.ndarray], optional): [..., 2] offset to translate. Defaults to Non

Tokens are integers from the set {0, 1, ..., vocab_size - 1}

Tokens are integers from the set {0, 1, ..., vocab_size - 1}