Create list of line segments from x and y coordinates, in the correct format for LineCollection: an array of the form numlines x (points per

first generate the building blocks, attn networks, encoders, decoders, PEs and Feedforward nets

Pad neighboring agent's history to the same length as that of the ego using NaNs

Reduce tensor at a dimension by named reduction functions. Args: x (torch.Tensor): tensor to be reduced reduction (str): one

scales `x` to [0, 1]

Returns the number of tokens used by a list of messages.

An optimization-based trajectory generator

shape: (5), batch_size, num_samp, num_agents, horizon, self.transition_dim

perturb the gradient and estimate the guidance loss w.r.t. the input trajectory Input: x_initial: [batch_size*num_samp, (

generate observation for the predicted scene f_step steps into the future Args: src_pos (torch.tensor[torch.float]): xy position in src

generate observation for every step of the predictions Args: data_batch (dict): input data dictionary pos_pred (torch.tensor[torc

INPUTS: - pos : (B, T, 2) - yaw : (B, T, 1) - speed : (B, T) - extent: (B, 3) - avail: (B, T) - norm_info : tuple of (add

Project the given points onto this Polyline. Args: xyh: Points to project, of shape (B, N, T, 3) points: lane points to project o

Rasterise agents for each scene in the batch at a given frame. :param state_index: the frame index in the scene :return: a dict mappi

Rasterise agents for each scene in the batch at a given frame. :param state_index: the frame index in the scene :return: a dict mappi

Get a frame from all scenes :param state_index: the frame index :return: a list of dict from EgoDatasets

Get a frame from all scenes :param state_index: the frame index :return: a list of dict from EgoDatasets

Paint agent histories onto an agent-centric map image

Paint agent histories onto an agent-centric map image

map raw src into features of dim 21+lane dim

Renders the semantic map at given x,y coordinates. Args: center_in_world (np.ndarray): XY of the image center in world ref system

Render ego-centric view, possibly with a location heatmap (if using SpatialPlanner)

Reparameterize for the backpropagation of z instead of q. This makes it so that we can backpropagate through the sampling of z from our e

Reset the previous simulation episode. Randomly sample a batch of new scenes unless specified in @scene_indices Args: sc

recursively get attributes

Rollout an environment for a number of episodes Args: env (BaseEnv): a base simulation environment (gym-like) policy (Rollout

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

Filter agents from zarr INPUT_FOLDER according to multiple thresholds and store a boolean array of the same shape.

:param agents_emb: (T_all, B, M, d_k) :param agent_masks: (B, T_obs, M) :return: (T_all, B, M, d_k)

Resets the test-time hard constraints to follow during prediction.

Instantiates test-time hard constraints using the config (dict) passed in.

Resets the test-time guidance functions to follow during prediction.

Resets the test-time guidance functions to follow during prediction.

Always([in stop sign region] -> Eventually( Always_[0, 3](low speed And in stop region) ) <=> Always( Not(in stop region) Or Eventual

:param agents_emb: (T_obs, B, M, d_k) :param agent_masks: (B, T_obs, M) :return: (T_obs, B, M, d_k)

:param agents_emb: (T_all, B, M, d_k) :param agent_masks: (B, T_hist, M), T_all = 1 when mlp is used :param aux_info: dict

extract positional info from src for all datatypes, e.g., for vehicles, the first two dimensions are x and y

extract positional info from src for all datatypes, e.g., for vehicles, the first two dimensions are x and y

Step the simulation with control inputs Args: actions (RolloutAction): action for controlling ego and/or agents

:param agents_emb: (T_fut, B, M, d_k) :param context: (T_obs, B, M, d_k) :param agent_masks: (B, T_obs, M) :return: (

:param agents_emb: (T_obs, B, M, d_k) :param agent_masks: (B, T_obs, M) :return: (T_obs, B, M, d_k)

Turning a trajectory tree in tiled form to a tree data structure Args: total_traj (torch.tensor or np.ndarray): tiled trajectory tree

Sends all torch tensors in nested dictionary or list or tuple to device @device, and returns a new nested structure. Args: x (di