Pad a nested dictionary or list or tuple of sequence tensors in the time dimension (dimension 1). Args: seq (dict or list or tuple):

Draw data samples (x') given a batch of condition inputs (x_c) and the VAE prior. Args: condition_inputs (dict, torch.Te

Get the dynamic states of the current timestep Note: Support both agent-centric and scene-centric (extra dimension for th

- query : query for GPT API.

Set the policy to evaluation mode

Recursively apply functions to a nested dictionary or list or tuple, given a dictionary of {data_type: function_to_apply}. Args:

Get checkpoint and config path given either a ngc job ID or a local dir. If a @ngc_job_id is specified, the function will first look for a d

Transform a set of 2D/3D points using the given transformation matrix. Assumes row major ordering of the input points. The transform function

Apply function @func to np.ndarray objects in a nested dictionary or list or tuple. Args: x (dict or list or tuple): a possibly

Produce N identical layers.

Args: avails (torch.Tensor): target availabilities [B, (A), T] Returns: last_indices (torch.Tensor): index o

Given the observation object, add Gaussian noise to positions and yaws Args: obs(Dict[torch.tensor]): observation dict R

Aggregated per-step loss between gt and predicted trajectories Args: predictions (torch.Tensor): predicted trajectory [B, (A), T, D]

Apply dynamics on input action trajectory to get state+action trajectory Note: Support both agent-centric and scene-centr

Evaluates all guidance losses and total and individual values. - x_loss: (B, N, T, 6) the trajectory to use to compute losses and 6 i

scale back the trajectory from standard normal distribution to original scale Note: Support both agent-centric and scene-

Queries the sim scene for the future state traj (in global frame) of all agents. - sim_scene : to query - fut_sec : how far in the future

Function to compute output shape from inputs to this module. Args: input_shape (iterable of int): shape of input. Does no

Input: - pos : (B, M, (Q), T, 2) - yaw : (B, M, (Q), T, 1) - speed : (B, M, (Q), T) - extent: (B, M, (Q), 3)

mark the image with manhattan distance to the drivable area Args: road_flag (torch.Tensor[B,W,H]): an image with 1 channel, 1 for drivabl

Apply dynamics on input action trajectory to get state+action trajectory Input: x_out: (batch_size, num_steps, 2). scaled

- traj: B x T x D

Aggregate per-step metrics for each scene. 1. if there are more than one agent per scene, aggregate their metrics for each scene using @agg_

A basic Conditional Variational Autoencoder Network (C-VAE) Args: q_net (nn.Module): a model that encodes data (x) and c

Check the shapes of args required by metrics Args: ground_truth (np.ndarray): array of shape (batch)x(timesteps)x(2D coords)

Infer vel and action (acc, yawvel) from state (x, y, yaw) based on Unicycle. Note: Support both agent-centric and scene-centric (extr

Apply function @func to torch.Tensor objects in a nested dictionary or list or tuple. Args: x (dict or list or tuple): a possibl

Converts trajdata's state format to pos, yaw, and speed. Set Nans to 0s

Converts local agent poses to global frame. Input: pos_pred: (num_agents, num_samp, time_steps, 2) yaw_pred: (num_agents, num

This will convert a scene-centric data batch to agent-centric coordinates keep_order_of_neighbors: if True, the order of neighbors (includin

scene_index: dict( action_positions=[num_agent, T, ...], action_yaws=[num_agent, T, ...],

Final step loss between gt and predicted trajectories (normally used in conjunction with a forward dynamics model) Args: predictions

Get indices of elements for which the bounding box described by bounds intersects the one defined around center (square with side 2*half_side