-param x: (B*N, M, T, transition_dim) -return: aux_info with 'neighbor_fut_feat': (B*N, M, M, T_fut, K_neigh)

Differs from parent only in the return type. Instead of doing everything from scratch, we rely on super call and fix the agents_mask

Cost for road departure.

Get per-scene rollout trajectory in the world coordinate system

Joins all dimensions between dimensions (@begin_axis, @end_axis) into a flat dimension, for all tensors in nested dictionary or list or tuple

Compute the log(sum(exp(x), dim)) in a numerically stable manner. Adapted from CS 236 at Stanford. Args: x (torch.Tensor): a tens

Pad input tensor or array @seq in the time dimension (dimension 1). Args: seq (np.ndarray or torch.Tensor): sequence to be padded

video_diffusion only perturb the gradient and estimate the guidance loss w.r.t. the input trajectory Input: x_ini

if self.predict_epsilon, model output is (scaled) noise; otherwise, model predicts x0 directly

if self.predict_epsilon, model output is (scaled) noise; otherwise, model predicts x0 directly

Bilinearly interpolates given positions in feature grid. - pos : (B x T x 2) float - feat_grid : (B x C x H x W) Returns: - (B x

-params neighbor_fut : (B*N/B, M, M, T_fut, K_neigh) -output neighbor_fut_feat : (B*N/B, M, T_fut, K_d)

Remove comment lines from a string containing Python source code.

Take samples using given a batch of distribution parameters, e.g., mean and logvar of a unit Gaussian Args: params (dict

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

select a slice of the tensor Args: x (torch.Tensor or np.array): the tensor dim (int): dimension to select start_idx (int

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

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

Translate a tbsim config to a l5kit config Args: cfg (ExperimentConfig): an ExperimentConfig instance Returns: cfg for

Translate a unified passthrough config to trajdata.

Integrate the state forward with initial state x0, action u Note: Support both agent-centric and scene-centric (extra dimension for t

Update any persistant state needed by guidance loss functions. - global_t : the current global timestep of rollout

Given a hard constraint config dict, verifies it's structured as expected. Should contain fields 'agents', 'loc', and 'times'

put the lists of ROI align result into embedding tensor with the help of index

KL divergence loss. Computes D_KL( N(mu, sigma) || N(0, 1) ). Note that this function averages across the batch dimension, but sums across di

KL divergence loss between two Gaussian distributions. This function computes the average loss across the batch. Args: mu_1 (torc

Given feature layers and proposals return bilinear interpolated points in feature layer Args: features (torch.Tensor): Tensor of shap

Get a PyTorch dataset object that can be used to train DNNs with vectorized input Args: cfg (dict): configuration file

Creates networks and places them into @self.nets.

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