SYMBOL INDEX (500 symbols across 54 files)

FILE: data_processing/nuplan/generate_nuplan_dataset.py
  function find_feature_paths (line 13) | def find_feature_paths(root_path, feature_name):
  function main (line 28) | def main(cfg):

FILE: data_processing/nuplan/preprocess_dataset_nuplan.py
  function _work_one_chunk (line 13) | def _work_one_chunk(idx, cfg_dict):
  function _run_one_cfg (line 22) | def _run_one_cfg(cfg):
  function main (line 40) | def main(cfg):

FILE: data_processing/postprocess_simulation_environments.py
  function main (line 12) | def main(cfg):

FILE: data_processing/waymo/add_nocturne_compatible_val_scenarios_to_test.py
  function add_val_to_test (line 10) | def add_val_to_test(cfg):
  function main (line 31) | def main(cfg):

FILE: data_processing/waymo/convert_pickles_to_jsons.py
  function reverse_ag_type_mapping (line 12) | def reverse_ag_type_mapping(agent_type_onehot):
  function compute_route_road_edges (line 19) | def compute_route_road_edges(route):
  function main (line 40) | def main(cfg):

FILE: data_processing/waymo/create_gpudrive_pickles.py
  function _mp_init (line 16) | def _mp_init():
  function _work_one_chunk (line 32) | def _work_one_chunk(idx, cfg_dict):
  function _run_one_cfg (line 41) | def _run_one_cfg(cfg):
  function main (line 73) | def main(cfg):

FILE: data_processing/waymo/create_waymo_eval_set.py
  function main (line 9) | def main(cfg):

FILE: data_processing/waymo/generate_k_disks_vocabulary.py
  function main (line 18) | def main(cfg):

FILE: data_processing/waymo/generate_waymo_dataset.py
  function poly_gon_and_line (line 27) | def poly_gon_and_line(poly_dict):
  function road_info_except_lane (line 42) | def road_info_except_lane(x_list, road_keys):
  function road_info_lane (line 66) | def road_info_lane(x_dict):
  function get_lane_pairs (line 86) | def get_lane_pairs(engage_lanes):
  function get_engage_lanes (line 133) | def get_engage_lanes(data):
  function get_lane_graph (line 149) | def get_lane_graph(data):
  function process_lanegraph (line 162) | def process_lanegraph(data):
  function _parse_object_state (line 218) | def _parse_object_state(states, final_state):
  function _init_object (line 244) | def _init_object(track):
  function get_objects (line 269) | def get_objects(scenario_list, index):
  function get_road_info (line 289) | def get_road_info(scenario_list, index):
  function collect_data (line 319) | def collect_data(cfg, output_path, files_path, files, chunk):
  function _work_one_chunk (line 386) | def _work_one_chunk(idx, cfg_dict):
  function _run_one_cfg (line 394) | def _run_one_cfg(cfg):
  function main (line 418) | def main(cfg):

FILE: data_processing/waymo/preprocess_dataset_waymo.py
  function _mp_init (line 17) | def _mp_init():
  function _work_one_chunk (line 33) | def _work_one_chunk(idx, cfg_dict):
  function _run_one_cfg (line 42) | def _run_one_cfg(cfg):
  function main (line 64) | def main(cfg):

FILE: datamodules/nuplan/nuplan_datamodule_autoencoder.py
  function worker_init_fn (line 8) | def worker_init_fn(worker_id):
  class NuplanDataModuleAutoEncoder (line 11) | class NuplanDataModuleAutoEncoder(pl.LightningDataModule):
    method __init__ (line 13) | def __init__(self,
    method setup (line 29) | def setup(self, stage):
    method train_dataloader (line 34) | def train_dataloader(self):
    method val_dataloader (line 44) | def val_dataloader(self):

FILE: datamodules/nuplan/nuplan_datamodule_ldm.py
  function worker_init_fn (line 7) | def worker_init_fn(worker_id):
  class NuplanDataModuleLDM (line 10) | class NuplanDataModuleLDM(pl.LightningDataModule):
    method __init__ (line 12) | def __init__(self,
    method setup (line 28) | def setup(self, stage):
    method train_dataloader (line 33) | def train_dataloader(self):
    method val_dataloader (line 43) | def val_dataloader(self):

FILE: datamodules/waymo/waymo_datamodule_autoencoder.py
  function worker_init_fn (line 8) | def worker_init_fn(worker_id):
  class WaymoDataModuleAutoEncoder (line 11) | class WaymoDataModuleAutoEncoder(pl.LightningDataModule):
    method __init__ (line 13) | def __init__(self,
    method setup (line 29) | def setup(self, stage):
    method train_dataloader (line 34) | def train_dataloader(self):
    method val_dataloader (line 44) | def val_dataloader(self):

FILE: datamodules/waymo/waymo_datamodule_ctrl_sim.py
  function worker_init_fn (line 8) | def worker_init_fn(worker_id):
  class WaymoDataModuleCtRLSim (line 11) | class WaymoDataModuleCtRLSim(pl.LightningDataModule):
    method __init__ (line 13) | def __init__(self,
    method setup (line 29) | def setup(self, stage):
    method train_dataloader (line 34) | def train_dataloader(self):
    method val_dataloader (line 44) | def val_dataloader(self):

FILE: datamodules/waymo/waymo_datamodule_ldm.py
  function worker_init_fn (line 7) | def worker_init_fn(worker_id):
  class WaymoDataModuleLDM (line 10) | class WaymoDataModuleLDM(pl.LightningDataModule):
    method __init__ (line 12) | def __init__(self,
    method setup (line 28) | def setup(self, stage):
    method train_dataloader (line 33) | def train_dataloader(self):
    method val_dataloader (line 43) | def val_dataloader(self):

FILE: datasets/nuplan/dataset_autoencoder_nuplan.py
  class NuplanDatasetAutoEncoder (line 26) | class NuplanDatasetAutoEncoder(Dataset):
    method __init__ (line 34) | def __init__(self, cfg: Any, split_name: str = "train", mode: str = "t...
    method get_lane_graph_within_fov (line 68) | def get_lane_graph_within_fov(self, lane_graph: Dict[str, Any]) -> Dic...
    method partition_compact_lane_graph (line 145) | def partition_compact_lane_graph(self, compact_lane_graph: Dict[str, A...
    method extract_lane_graph (line 237) | def extract_lane_graph(
    method extract_agents (line 306) | def extract_agents(self, ego, vehicles, pedestrians, static_objects):
    method get_agents_within_fov (line 397) | def get_agents_within_fov(self, agent_states, agent_types):
    method get_road_points_adj (line 424) | def get_road_points_adj(self, compact_lane_graph):
    method get_partitioned_masks (line 478) | def get_partitioned_masks(self, agents, lanes, a2a_edge_index, l2l_edg...
    method get_data (line 503) | def get_data(self, data, idx):
    method get (line 716) | def get(self, idx: int):
    method len (line 740) | def len(self):
  function main (line 744) | def main(cfg):

FILE: datasets/nuplan/dataset_ldm_nuplan.py
  class NuplanDatasetLDM (line 24) | class NuplanDatasetLDM(Dataset):
    method __init__ (line 25) | def __init__(self, cfg: Any, split_name: str = "train") -> None:
    method get_data (line 47) | def get_data(self, data, idx):
    method get (line 106) | def get(self, idx: int):
    method len (line 117) | def len(self):
  function main (line 121) | def main(cfg):

FILE: datasets/waymo/dataset_autoencoder_waymo.py
  class WaymoDatasetAutoEncoder (line 32) | class WaymoDatasetAutoEncoder(Dataset):
    method __init__ (line 41) | def __init__(self, cfg: Any, split_name: str = "train", mode: str = "t...
    method partition_compact_lane_graph (line 95) | def partition_compact_lane_graph(self, compact_lane_graph: Dict[str, A...
    method normalize_compact_lane_graph (line 201) | def normalize_compact_lane_graph(self, lane_graph: Dict[str, Any], nor...
    method get_lane_graph_within_fov (line 237) | def get_lane_graph_within_fov(self, lane_graph: Dict[str, Any]) -> Dic...
    method get_road_points_adj (line 330) | def get_road_points_adj(
    method get_agents_within_fov (line 416) | def get_agents_within_fov(
    method remove_offroad_agents (line 477) | def remove_offroad_agents(
    method get_partitioned_masks (line 527) | def get_partitioned_masks(
    method get_data (line 587) | def get_data(self,
    method get (line 997) | def get(self, idx: int):
    method len (line 1013) | def len(self):
  function main (line 1017) | def main(cfg):

FILE: datasets/waymo/dataset_ctrl_sim.py
  class CtRLSimDataset (line 30) | class CtRLSimDataset(Dataset):
    method __init__ (line 43) | def __init__(self, cfg, split_name='train'):
    method get_upsampled_and_sd_lanes (line 95) | def get_upsampled_and_sd_lanes(self, compact_lane_graph):
    method remove_offroad_agents (line 112) | def remove_offroad_agents(self, agent_states, agent_types, lanes):
    method rollout_k_disks (line 135) | def rollout_k_disks(self, agent_states):
    method get_ego_collision_rewards (line 246) | def get_ego_collision_rewards(self, agent_states_all):
    method get_last_valid_positions (line 279) | def get_last_valid_positions(self, states):
    method get_agent_mask (line 290) | def get_agent_mask(self, agent_states, normalize_dict, fov=None):
    method compute_rtgs (line 306) | def compute_rtgs(self, rewards):
    method select_closest_max_num_agents (line 327) | def select_closest_max_num_agents(
    method get_normalized_lanes_in_fov (line 394) | def get_normalized_lanes_in_fov(self, lanes, normalize_dict):
    method discretize_rtgs (line 433) | def discretize_rtgs(self, rtgs):
    method collect_state_transitions (line 441) | def collect_state_transitions(self, data):
    method get_data (line 474) | def get_data(self, data, idx):
    method get (line 709) | def get(self, idx):
    method len (line 745) | def len(self):
  function main (line 750) | def main(cfg):

FILE: datasets/waymo/dataset_ldm_waymo.py
  class WaymoDatasetLDM (line 23) | class WaymoDatasetLDM(Dataset):
    method __init__ (line 24) | def __init__(self, cfg: Any, split_name: str = "train") -> None:
    method get_data (line 46) | def get_data(self, data, idx):
    method get (line 105) | def get(self, idx: int):
    method len (line 116) | def len(self):
  function main (line 120) | def main(cfg):

FILE: eval.py
  function generate_simulation_environments (line 18) | def generate_simulation_environments(cfg, cfg_ae, save_dir=None):
  function eval_ldm (line 42) | def eval_ldm(cfg, cfg_ae, save_dir=None):
  function eval_autoencoder (line 79) | def eval_autoencoder(cfg, save_dir=None):
  function main (line 105) | def main(cfg):

FILE: metrics.py
  class Metrics (line 7) | class Metrics():
    method __init__ (line 9) | def __init__(self, cfg):
    method compute_metrics (line 16) | def compute_metrics(self):

FILE: models/ctrl_sim.py
  class CtRLSim (line 11) | class CtRLSim(pl.LightningModule):
    method __init__ (line 13) | def __init__(self, cfg):
    method forward (line 28) | def forward(self, data, eval=False):
    method compute_loss (line 36) | def compute_loss(self, data, preds):
    method training_step (line 99) | def training_step(self, data, batch_idx):
    method validation_step (line 138) | def validation_step(self, data, batch_idx):
    method on_before_optimizer_step (line 165) | def on_before_optimizer_step(self, optimizer):
    method configure_optimizers (line 175) | def configure_optimizers(self):

FILE: models/scenario_dreamer_autoencoder.py
  function worker_init_fn (line 22) | def worker_init_fn(worker_id):
  class ScenarioDreamerAutoEncoder (line 26) | class ScenarioDreamerAutoEncoder(pl.LightningModule):
    method __init__ (line 28) | def __init__(self, cfg):
    method test_dataloader (line 45) | def test_dataloader(self):
    method forward (line 71) | def forward(self, data):
    method _log_losses (line 92) | def _log_losses(self, loss_dict, split='train', batch_size=None):
    method _cache_latents (line 118) | def _cache_latents(self, data):
    method training_step (line 187) | def training_step(self, data, batch_idx):
    method validation_step (line 195) | def validation_step(self, data, batch_idx):
    method test_step (line 229) | def test_step(self, data, batch_idx):
    method on_before_optimizer_step (line 265) | def on_before_optimizer_step(self, optimizer):
    method configure_optimizers (line 273) | def configure_optimizers(self):

FILE: models/scenario_dreamer_ldm.py
  class ScenarioDreamerLDM (line 27) | class ScenarioDreamerLDM(pl.LightningModule):
    method __init__ (line 28) | def __init__(self, cfg, cfg_ae):
    method on_train_start (line 42) | def on_train_start(self):
    method optimizer_step (line 47) | def optimizer_step(self, *args, **kwargs):
    method _log_losses (line 53) | def _log_losses(self, loss_dict, split='train', batch_size=None):
    method training_step (line 79) | def training_step(self, data, batch_idx):
    method validation_step (line 87) | def validation_step(self, data, batch_idx):
    method forward (line 118) | def forward(self,
    method _build_ldm_dset_from_ae_dset_for_inpainting (line 191) | def _build_ldm_dset_from_ae_dset_for_inpainting(self, ae_dset, batch_s...
    method _initialize_pyg_dset (line 294) | def _initialize_pyg_dset(self, mode, num_samples, batch_size, conditio...
    method generate (line 448) | def generate(
    method on_before_optimizer_step (line 515) | def on_before_optimizer_step(self, optimizer):
    method on_save_checkpoint (line 522) | def on_save_checkpoint(self, checkpoint):
    method on_load_checkpoint (line 527) | def on_load_checkpoint(self, checkpoint):
    method configure_optimizers (line 533) | def configure_optimizers(self):

FILE: nn_modules/autoencoder.py
  class ScenarioDreamerEncoder (line 15) | class ScenarioDreamerEncoder(nn.Module):
    method __init__ (line 18) | def __init__(self, cfg):
    method forward (line 75) | def forward(
  class ScenarioDreamerDecoder (line 167) | class ScenarioDreamerDecoder(nn.Module):
    method __init__ (line 170) | def __init__(self, cfg):
    method forward (line 221) | def forward(
  class AutoEncoder (line 301) | class AutoEncoder(nn.Module):
    method __init__ (line 304) | def __init__(self, cfg):
    method loss (line 321) | def loss(self, data):
    method forward_encoder (line 404) | def forward_encoder(self, data, return_stats=False, return_lane_embedd...
    method forward_decoder (line 435) | def forward_decoder(self, agent_latents, lane_latents, data):
    method forward (line 448) | def forward(self, data, return_latents=False, return_lane_embeddings=F...

FILE: nn_modules/ctrl_sim.py
  class CtRLSimMapEncoder (line 8) | class CtRLSimMapEncoder(nn.Module):
    method __init__ (line 10) | def __init__(self, cfg):
    method get_road_pts_mask (line 38) | def get_road_pts_mask(self, roads):
    method forward (line 45) | def forward(self, data):
  class CtRLSimEncoder (line 83) | class CtRLSimEncoder(nn.Module):
    method __init__ (line 85) | def __init__(self, cfg):
    method forward (line 129) | def forward(self, data, eval):
  class CtRLSimDecoder (line 249) | class CtRLSimDecoder(nn.Module):
    method __init__ (line 251) | def __init__(self, cfg):
    method forward (line 291) | def forward(self, data, scene_enc, eval=False):

FILE: nn_modules/dit.py
  class DiT (line 9) | class DiT(nn.Module):
    method __init__ (line 11) | def __init__(self, cfg):
    method initialize_weights (line 58) | def initialize_weights(self):
    method forward (line 109) | def forward(self,

FILE: nn_modules/ldm.py
  class LDM (line 12) | class LDM(nn.Module):
    method __init__ (line 13) | def __init__(self, cfg):
    method predict_start_from_noise (line 59) | def predict_start_from_noise(self, x_t, t, noise):
    method q_posterior (line 67) | def q_posterior(self, x_start, x_t, t):
    method p_mean_variance (line 77) | def p_mean_variance(self, x_agent, x_lane, data, t_agent, t_lane):
    method p_sample (line 101) | def p_sample(self, x_agent, x_lane, data, t_agent, t_lane):
    method p_sample_loop (line 127) | def p_sample_loop(
    method forward (line 202) | def forward(self, data, mode='initial_scene'):
    method q_sample (line 217) | def q_sample(self, x_start, t, noise=None):
    method p_losses (line 227) | def p_losses(
    method loss (line 267) | def loss(self, data):

FILE: policies/idm_policy.py
  class IDMPolicy (line 17) | class IDMPolicy:
    method __init__ (line 29) | def __init__(self, cfg, env):
    method reset (line 57) | def reset(self, obs):
    method update_running_statistics (line 66) | def update_running_statistics(self, data_dict, scenario_dict, scene_co...
    method compute_metrics (line 152) | def compute_metrics(self):
    method act (line 200) | def act(self, obs, is_planner=True):
    method select_action (line 204) | def select_action(self, obs, is_planner=True):
    method _plot_lanes (line 247) | def _plot_lanes(self, lane_dict, agent_states, filename="lane_graph", ...
    method _plot_ego_path (line 290) | def _plot_ego_path(self, ego_path_polygon, agent_occupancies, ego_id):
    method _get_lane_data (line 309) | def _get_lane_data(self):
    method _get_closest_lane_point_from_position (line 330) | def _get_closest_lane_point_from_position(self, lane_points, position):
    method _get_closest_lane_from_position (line 336) | def _get_closest_lane_from_position(self, position, lane_geometries, m...
    method _compute_all_agent_lanes (line 350) | def _compute_all_agent_lanes(self, agent_states):
    method _get_ego_path (line 374) | def _get_ego_path(self, agent_id):
    method _initialize_ego_path (line 382) | def _initialize_ego_path(self, actor_id):
    method _compute_agent_occupancies (line 447) | def _compute_agent_occupancies(self, agent_states):
    method _compute_leading_agents_occ (line 464) | def _compute_leading_agents_occ(self, agent_states):
    method _get_accelerations (line 547) | def _get_accelerations(self, agent_states):
    method _get_steerings (line 598) | def _get_steerings(self, agent_states):
    method _get_next_states (line 648) | def _get_next_states(self, agent_states):
    method _get_next_path_position (line 716) | def _get_next_path_position(self, ego_path, current_lane, current_lane...

FILE: policies/rl_policy.py
  class RLPolicy (line 3) | class RLPolicy:
    method __init__ (line 4) | def __init__(self, cfg):
    method act (line 14) | def act(self, obs):

FILE: run_simulation.py
  class PolicyEvaluator (line 13) | class PolicyEvaluator:
    method __init__ (line 15) | def __init__(self, cfg, policy, env):
    method reset (line 23) | def reset(self):
    method update_running_statistics (line 35) | def update_running_statistics(self, info):
    method compute_metrics (line 43) | def compute_metrics(self):
    method evaluate_policy (line 55) | def evaluate_policy(self):
  function main (line 98) | def main(cfg):

FILE: simulator.py
  class Simulator (line 40) | class Simulator:
    method __init__ (line 49) | def __init__(self, cfg):
    method load_initial_scene (line 78) | def load_initial_scene(self, i):
    method _find_invalid_new_agents (line 106) | def _find_invalid_new_agents(
    method step (line 174) | def step(self, action):
    method _get_observation (line 350) | def _get_observation(self):
    method _update_viz_state (line 399) | def _update_viz_state(self, num_route_points=30):
    method initialize_data_dict (line 433) | def initialize_data_dict(self):
    method reset (line 476) | def reset(self, i):
    method render_state (line 536) | def render_state(self, name, movie_path):
  class CtRLSimBehaviourModel (line 577) | class CtRLSimBehaviourModel:
    method __init__ (line 582) | def __init__(self,
    method update_running_statistics (line 627) | def update_running_statistics(
    method compute_metrics (line 819) | def compute_metrics(self):
    method reset (line 844) | def reset(self, num_agents):
    method update_state (line 853) | def update_state(self, data_dict):
    method get_motion_data (line 878) | def get_motion_data(self, data_dict):
    method get_tilt_logits (line 1030) | def get_tilt_logits(self, tilt):
    method process_predicted_rtg (line 1037) | def process_predicted_rtg(
    method predict (line 1074) | def predict(self, motion_datas, data_dict, correspondences):
    method step (line 1129) | def step(self, data_dict):

FILE: train.py
  function train_ctrl_sim (line 21) | def train_ctrl_sim(cfg, save_dir=None):
  function train_ldm (line 70) | def train_ldm(cfg, cfg_ae, save_dir=None):
  function train_autoencoder (line 134) | def train_autoencoder(cfg, save_dir=None):
  function main (line 186) | def main(cfg):

FILE: utils/collision_helpers.py
  function compute_corners (line 4) | def compute_corners(positions, headings, lengths, widths):
  function get_axes (line 36) | def get_axes(vertices):
  function is_colliding (line 50) | def is_colliding(poly1, poly2):
  function batched_collision_checker (line 77) | def batched_collision_checker(ego_state, agent_states):
  function compute_collision_states_one_scene (line 112) | def compute_collision_states_one_scene(vehicles):

FILE: utils/data_container.py
  class CtRLSimData (line 7) | class CtRLSimData(HeteroData):
    method __inc__ (line 9) | def __inc__(self, key, value, store):
  class ScenarioDreamerData (line 13) | class ScenarioDreamerData(HeteroData):
    method __inc__ (line 15) | def __inc__(self, key, value, store):
  function get_batches (line 29) | def get_batches(data):
  function get_features (line 39) | def get_features(data):
  function get_edge_indices (line 65) | def get_edge_indices(data):
  function get_encoder_edge_indices (line 80) | def get_encoder_edge_indices(data):

FILE: utils/data_helpers.py
  function extract_raw_waymo_data (line 10) | def extract_raw_waymo_data(agents_data: List[Dict[str, Any]]) -> Tuple[n...
  function add_batch_dim (line 68) | def add_batch_dim(arr):
  function get_object_type_onehot_waymo (line 71) | def get_object_type_onehot_waymo(agent_type):
  function get_lane_connection_type_onehot_waymo (line 76) | def get_lane_connection_type_onehot_waymo(lane_connection_type):
  function get_lane_connection_type_onehot_nuplan (line 81) | def get_lane_connection_type_onehot_nuplan(lane_connection_type):
  function get_lane_type_onehot_nuplan (line 86) | def get_lane_type_onehot_nuplan(lane_type):
  function get_object_type_onehot_nuplan (line 91) | def get_object_type_onehot_nuplan(agent_type):
  function reorder_indices (line 97) | def reorder_indices(
  function modify_agent_states (line 312) | def modify_agent_states(agent_states):
  function normalize_scene (line 337) | def normalize_scene(
  function unnormalize_scene (line 379) | def unnormalize_scene(
  function randomize_indices (line 421) | def randomize_indices(
  function normalize_latents (line 455) | def normalize_latents(
  function unnormalize_latents (line 469) | def unnormalize_latents(
  function reparameterize (line 483) | def reparameterize(mu, log_var):
  function sample_latents (line 497) | def sample_latents(
  function convert_batch_to_scenarios (line 525) | def convert_batch_to_scenarios(data, batch_size, batch_idx, cache_dir, c...

FILE: utils/diffusion_helpers.py
  function extract (line 4) | def extract(a, t, x_shape):
  function cosine_beta_schedule (line 10) | def cosine_beta_schedule(timesteps, s=0.008, dtype=torch.float32):

FILE: utils/dit_layers.py
  function modulate (line 11) | def modulate(x, shift, scale):
  function _ntuple (line 15) | def _ntuple(n):
  function get_1d_sincos_pos_embed_from_grid (line 29) | def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
  class AttentionLayerDiT (line 51) | class AttentionLayerDiT(MessagePassing):
    method __init__ (line 53) | def __init__(self,
    method message (line 82) | def message(self, q_i, k_j, v_j, index, ptr):
    method update (line 88) | def update(self, inputs):
    method _attn_block (line 92) | def _attn_block(self, x_src, x_dst, edge_index):
    method forward (line 104) | def forward(self, x, edge_index):
  class Mlp (line 109) | class Mlp(nn.Module):
    method __init__ (line 113) | def __init__(
    method forward (line 138) | def forward(self, x):
  class LabelEmbedder (line 148) | class LabelEmbedder(nn.Module):
    method __init__ (line 153) | def __init__(self, num_classes, hidden_size, dropout_prob):
    method token_drop (line 160) | def token_drop(self, labels, force_drop_ids=None):
    method forward (line 171) | def forward(self, labels, train, force_drop_ids=None):
  class TimestepEmbedder (line 180) | class TimestepEmbedder(nn.Module):
    method __init__ (line 185) | def __init__(self, hidden_size, frequency_embedding_size=256):
    method timestep_embedding (line 195) | def timestep_embedding(t, dim, max_period=10000):
    method forward (line 215) | def forward(self, t):
  class DiTBlock (line 221) | class DiTBlock(nn.Module):
    method __init__ (line 226) | def __init__(self, hidden_size, num_heads, dropout, mlp_ratio=4.0, **b...
    method forward (line 239) | def forward(self, x, c, edge_index):
  class FactorizedDiTBlock (line 248) | class FactorizedDiTBlock(nn.Module):
    method __init__ (line 252) | def __init__(
    method forward (line 284) | def forward(
  class FinalLayer (line 315) | class FinalLayer(nn.Module):
    method __init__ (line 320) | def __init__(self, hidden_size, latent_size):
    method forward (line 329) | def forward(self, x, c):
  class TwoLayerResMLP (line 336) | class TwoLayerResMLP(nn.Module):
    method __init__ (line 337) | def __init__(self, input_dim, hidden_dim):
    method forward (line 349) | def forward(self, x):

FILE: utils/geometry.py
  function rotate_and_normalize_angles (line 3) | def rotate_and_normalize_angles(current_angles, rotation_angle):
  function normalize_angle (line 11) | def normalize_angle(angle):
  function dot_product_2d (line 15) | def dot_product_2d(a, b):
  function cross_product_2d (line 20) | def cross_product_2d(a, b):
  function make_2d_rotation_matrix (line 24) | def make_2d_rotation_matrix(angle_in_radians):
  function apply_se2_transform (line 30) | def apply_se2_transform(coordinates, translation, yaw):
  function radians_to_degrees (line 43) | def radians_to_degrees(radians):
  function normalize_lanes_and_agents (line 48) | def normalize_lanes_and_agents(agents, lanes, normalize_dict, dataset):
  function normalize_agents (line 79) | def normalize_agents(agent_states, normalize_dict, offset=np.pi/2):
  function normalize_lanes (line 102) | def normalize_lanes(lanes, normalize_dict, offset=np.pi/2):

FILE: utils/gpudrive_helpers.py
  class EntityType (line 58) | class EntityType(Enum):
  class MapType (line 75) | class MapType(Enum):
  class MapVector2 (line 103) | class MapVector2:
    method __init__ (line 105) | def __init__(self, x: float = 0.0, y: float = 0.0):
  class MapVector3 (line 110) | class MapVector3:
    method __init__ (line 112) | def __init__(self, x: float = 0.0, y: float = 0.0, z: float = 0.0):
  class VehicleSize (line 119) | class VehicleSize:
    method __init__ (line 121) | def __init__(self, length: float = 0.0, width: float = 0.0, height: fl...
  class MetaData (line 128) | class MetaData:
    method __init__ (line 130) | def __init__(self):
  class MapObject (line 138) | class MapObject:
    method __init__ (line 140) | def __init__(self):
  class MapRoad (line 159) | class MapRoad:
    method __init__ (line 161) | def __init__(self):
  class Map (line 171) | class Map:
    method __init__ (line 173) | def __init__(self):
  function distance_2d (line 186) | def distance_2d(p1: MapVector2, p2: MapVector2) -> float:
  function get_ego_state (line 192) | def get_ego_state(ego_state):
  function get_partner_obs (line 242) | def get_partner_obs(
  function get_map_obs (line 348) | def get_map_obs(
  function get_route_obs (line 523) | def get_route_obs(
  class ForwardKinematics (line 624) | class ForwardKinematics:
    method __init__ (line 628) | def __init__(
    method forward_kinematics (line 646) | def forward_kinematics(self, action):
  function get_action_value_tensor (line 696) | def get_action_value_tensor() -> torch.Tensor:
  function _normalize_min_max (line 741) | def _normalize_min_max(tensor, min_val, max_val):
  function _angle_add (line 755) | def _angle_add(angle1: float, angle2: float) -> float:
  function from_json_MapVector2 (line 762) | def from_json_MapVector2(j: dict) -> MapVector2:
  function from_json_MapObject (line 773) | def from_json_MapObject(j: dict) -> MapObject:
  function from_json_MapRoad (line 856) | def from_json_MapRoad(j: dict, polylineReductionThreshold: float = 0.0) ...
  function calc_mean (line 992) | def calc_mean(j: dict) -> Tuple[float, float]:
  function make_road_edge (line 1035) | def make_road_edge(road_init, j, world_mean) -> dict:
  function create_road_edges (line 1132) | def create_road_edges(data, world_mean, max_num_edges=10000) -> List[dict]:
  function from_json_Map (line 1209) | def from_json_Map(j: dict, polylineReductionThreshold: float = 0.0) -> d...
  function load_policy (line 1436) | def load_policy(path_to_cpt, model_name, device):
  function log_prob (line 1462) | def log_prob(logits, value):
  function entropy (line 1470) | def entropy(logits):
  function sample_logits (line 1478) | def sample_logits(
  function layer_init (line 1519) | def layer_init(layer, std=np.sqrt(2), bias_const=0.0):
  class NeuralNet (line 1528) | class NeuralNet(
    method __init__ (line 1534) | def __init__(
    method encode_observations (line 1631) | def encode_observations(self, observation):
    method forward (line 1664) | def forward(self, obs, action=None, deterministic=False):
    method unpack_obs (line 1677) | def unpack_obs(self, obs_flat):

FILE: utils/inpainting_helpers.py
  function normalize_and_crop_scene (line 15) | def normalize_and_crop_scene(cond_d, new_d, normalize_dict, cfg, dataset...
  function sample_num_lanes_agents_inpainting (line 196) | def sample_num_lanes_agents_inpainting(

FILE: utils/k_disks_helpers.py
  function compute_k_disks_vocabulary (line 5) | def compute_k_disks_vocabulary(state_transitions, vocab_size, l, w, eps):
  function transform_box_corners_from_vocab (line 54) | def transform_box_corners_from_vocab(box_coords, V):
  function get_local_state_transition (line 84) | def get_local_state_transition(current_state, next_state):
  function transform_box_corners_from_local_state (line 106) | def transform_box_corners_from_local_state(box_coords, local_state_trans...
  function get_global_next_state (line 133) | def get_global_next_state(global_states, local_transitions):
  function forward_k_disks (line 162) | def forward_k_disks(states, actions, vocab, delta_t, exists):
  function inverse_k_disks (line 184) | def inverse_k_disks(states, next_states, vocab):

FILE: utils/lane_graph_helpers.py
  function get_compact_lane_graph (line 5) | def get_compact_lane_graph(data):
  function find_lane_groups (line 130) | def find_lane_groups(pre_pairs, suc_pairs):
  function find_lane_group_id (line 187) | def find_lane_group_id(lane_id, lane_groups):
  function resample_polyline (line 194) | def resample_polyline(points, num_points=20):
  function resample_polyline_every (line 213) | def resample_polyline_every(polyline, every=1.5):
  function resample_lanes (line 234) | def resample_lanes(lanes, num_points):
  function resample_lanes_with_mask (line 243) | def resample_lanes_with_mask(lanes, lanes_mask, num_points):
  function adjacency_matrix_to_adjacency_list (line 254) | def adjacency_matrix_to_adjacency_list(lane_graph_adj):
  function estimate_heading (line 272) | def estimate_heading(positions):
  function find_closest_lane (line 287) | def find_closest_lane(lanes, pos):

FILE: utils/layers.py
  class ResidualMLP (line 7) | class ResidualMLP(nn.Module):
    method __init__ (line 21) | def __init__(self, input_dim: int, hidden_dim: int,
    method forward (line 44) | def forward(self, x: torch.Tensor) -> torch.Tensor:
  class MLP (line 62) | class MLP(nn.Module):
    method __init__ (line 65) | def __init__(self, input_dim, hidden_dim, output_dim):
    method forward (line 75) | def forward(self, x):
  class AttentionLayer (line 79) | class AttentionLayer(MessagePassing):
    method __init__ (line 81) | def __init__(self,
    method forward (line 128) | def forward(self, x, r, edge_index):
    method message (line 142) | def message(self, q_i, k_j, v_j, r, index, ptr):
    method update (line 151) | def update(self, inputs, x_dst):
    method _attn_block (line 156) | def _attn_block(self, x_src, x_dst, r, edge_index):
    method _ff_block (line 163) | def _ff_block(self, x):
  class EdgeFeatureUpdate (line 167) | class EdgeFeatureUpdate(MessagePassing):
    method __init__ (line 169) | def __init__(self, node_hidden_dim, edge_hidden_dim):
    method forward (line 178) | def forward(self, x, edge_index, edge_attr):
    method edge_update (line 181) | def edge_update(self, x_i, x_j, edge_attr):
  class AutoEncoderFactorizedAttentionBlock (line 190) | class AutoEncoderFactorizedAttentionBlock(nn.Module):
    method __init__ (line 192) | def __init__(self,
    method forward (line 246) | def forward(self, agent_embeddings,

FILE: utils/losses.py
  class GeometricLoss (line 5) | class GeometricLoss(nn.Module):
    method __init__ (line 6) | def __init__(self, mean_dim=1, apply_mean=True):
    method forward (line 11) | def forward(self, pred, targ, batch):
  class GeometricCrossEntropy (line 31) | class GeometricCrossEntropy(GeometricLoss):
    method _loss (line 33) | def _loss(self, pred, targ):
  class GeometricHuber (line 36) | class GeometricHuber(GeometricLoss):
    method _loss (line 38) | def _loss(self, pred, targ):
  class GeometricL2 (line 41) | class GeometricL2(GeometricLoss):
    method _loss (line 43) | def _loss(self, pred, targ):
  class GeometricL1 (line 46) | class GeometricL1(GeometricLoss):
    method _loss (line 48) | def _loss(self, pred, targ):
  class GeometricKL (line 51) | class GeometricKL(GeometricLoss):
    method _loss (line 53) | def _loss(self, mu, log_var):

FILE: utils/metrics_helpers.py
  function compute_frechet_distance (line 12) | def compute_frechet_distance(X1, X2, apply_sqrt=True):
  function jsd (line 39) | def jsd(sim, gt, clip_min, clip_max, bin_size):
  function compute_vehicle_circles (line 57) | def compute_vehicle_circles(xy_position, heading, length, width):
  function compute_collision_rate (line 74) | def compute_collision_rate(samples):
  function get_compact_lane_graph (line 121) | def get_compact_lane_graph(G, lanes, num_points_per_lane=20):
  function _get_sledge_lane_graph_nuplan (line 184) | def _get_sledge_lane_graph_nuplan(data):
  function get_networkx_lane_graph_without_traffic_lights (line 303) | def get_networkx_lane_graph_without_traffic_lights(data):
  function get_networkx_lane_graph (line 350) | def get_networkx_lane_graph(data):
  function convert_data_to_unified_format (line 378) | def convert_data_to_unified_format(data, dataset_name):
  function get_lane_length (line 412) | def get_lane_length(positions):
  function compute_route_length (line 420) | def compute_route_length(samples):
  function compute_endpoint_dist (line 466) | def compute_endpoint_dist(samples):
  function get_keypoint_G (line 491) | def get_keypoint_G(G, lanes):
  function get_num_keypoints (line 537) | def get_num_keypoints(G):
  function get_degree_keypoints (line 542) | def get_degree_keypoints(G):
  function urban_planning_reach_and_convenience (line 549) | def urban_planning_reach_and_convenience(G_edges):
  function get_onroad_vehicles (line 566) | def get_onroad_vehicles(vehicles, lanes, tol=1.5):
  function get_nearest_dists (line 577) | def get_nearest_dists(vehicles):
  function get_lateral_devs (line 587) | def get_lateral_devs(vehicles, lanes):
  function get_angular_devs (line 597) | def get_angular_devs(vehicles, lanes):
  function get_lengths (line 627) | def get_lengths(vehicles):
  function get_widths (line 632) | def get_widths(vehicles):
  function get_speeds (line 637) | def get_speeds(vehicles):
  function compute_urban_planning_metrics (line 642) | def compute_urban_planning_metrics(samples, gt_samples):
  function compute_jsd_metrics (line 703) | def compute_jsd_metrics(samples, gt_samples):
  function compute_lane_metrics (line 774) | def compute_lane_metrics(samples, gt_samples):
  function compute_agent_metrics (line 793) | def compute_agent_metrics(samples, gt_samples):
  function compute_sim_agent_jsd_metrics (line 809) | def compute_sim_agent_jsd_metrics(

FILE: utils/pyg_helpers.py
  function get_edge_index_bipartite (line 3) | def get_edge_index_bipartite(num_src, num_dst):
  function get_edge_index_complete_graph (line 15) | def get_edge_index_complete_graph(graph_size):
  function get_indices_within_scene (line 23) | def get_indices_within_scene(batch):

FILE: utils/sim_env_helpers.py
  function postprocess_sim_env (line 23) | def postprocess_sim_env(
  function get_route_lane_indices (line 207) | def get_route_lane_indices(
  function clean_up_scene (line 301) | def clean_up_scene(data, dataset, mode='initial_scene', endpoint_thresho...
  function check_scene_validity (line 476) | def check_scene_validity(data, dataset):
  function check_scene_validity_inpainting (line 518) | def check_scene_validity_inpainting(data, dataset, heading_tolerance=np....
  function sample_route (line 590) | def sample_route(d, dataset, heading_tolerance=np.pi/3, num_points_in_ro...
  function get_default_route_center_yaw (line 654) | def get_default_route_center_yaw(dataset):
  function generate_simulation_environments (line 662) | def generate_simulation_environments(model, cfg, save_dir):
  function _transform_scene (line 851) | def _transform_scene(agents, lanes, route, transform_dict, dataset):
  function _extend_simulation_environment (line 885) | def _extend_simulation_environment(current_env, new_tile, target_route_l...
  function _sample_candidate (line 995) | def _sample_candidate(candidates, dataset):
  function _near_border (line 1017) | def _near_border(pos, fov=64, threshold=1):
  function _near_partition (line 1024) | def _near_partition(pos, dataset, threshold=2.5):
  function _valid_route_end (line 1030) | def _valid_route_end(lane_id, lane, fov=64, border_threshold=1, heading_...
  function _transform_corners (line 1055) | def _transform_corners(corners, transform_dict, dataset):
  function _check_overlapping_tiles (line 1071) | def _check_overlapping_tiles(new_tile_corners, existing_tiles, ignore_la...

FILE: utils/sim_helpers.py
  function ego_completed_route (line 9) | def ego_completed_route(ego_state, route, dist_threshold=2.0):
  function ego_collided (line 18) | def ego_collided(ego_state, agent_states, agent_scale=1.0):
  function ego_off_route (line 38) | def ego_off_route(ego_state, route, off_route_threshold=5.0):
  function ego_progress (line 49) | def ego_progress(ego_state, route):
  function normalize_route (line 60) | def normalize_route(route, normalize_dict, offset=np.pi/2):
  function get_ego_route (line 74) | def get_ego_route(compact_lane_graph, lanes, ego_trajectory, dist_thresh...

FILE: utils/sledge_helpers.py
  function calculate_progress (line 5) | def calculate_progress(path):
  function get_path_length (line 19) | def get_path_length(path):
  function interpolate_path (line 23) | def interpolate_path(distances, length, progress, states_se2_array, as_a...
  function coords_in_frame (line 40) | def coords_in_frame(coords, frame):
  function find_consecutive_true_indices (line 55) | def find_consecutive_true_indices(mask):
  function pixel_in_frame (line 74) | def pixel_in_frame(pixel, pixel_frame):
  function coords_to_pixel (line 90) | def coords_to_pixel(coords, frame, pixel_size):

FILE: utils/torch_helpers.py
  function from_numpy (line 4) | def from_numpy(data):

FILE: utils/train_helpers.py
  function create_lambda_lr_cosine (line 13) | def create_lambda_lr_cosine(cfg):
  function create_lambda_lr_linear (line 24) | def create_lambda_lr_linear(cfg):
  function create_lambda_lr_constant (line 35) | def create_lambda_lr_constant(cfg):
  function weight_init (line 43) | def weight_init(m):
  function cache_latent_stats (line 112) | def cache_latent_stats(cfg):
  function set_latent_stats (line 173) | def set_latent_stats(cfg):
  function get_causal_mask (line 191) | def get_causal_mask(cfg, num_timesteps, num_types):

FILE: utils/viz.py
  function plot_scene (line 12) | def plot_scene(
  function plot_lane_graph (line 183) | def plot_lane_graph(
  function visualize_batch (line 260) | def visualize_batch(num_samples,
  function plot_k_disks_vocabulary (line 349) | def plot_k_disks_vocabulary(V, png_path, dpi=1000):
  function render_state (line 362) | def render_state(
  function generate_video (line 518) | def generate_video(name, output_dir, delete_images=False):