Skip to content

Include simple multiprocessing #153

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Include simple multiprocessing #153

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ccd4a1f
Update run_create_submission_pickle.py
dXDb Sep 3, 2025
5f1b8ba
Update run_create_submission_pickle_challenge.py
dXDb Sep 3, 2025
7754238
Update run_create_submission_pickle.py
dXDb Sep 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 33 additions & 18 deletions navsim/planning/script/run_create_submission_pickle.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,26 @@
CONFIG_NAME = "default_run_create_submission_pickle"


from navsim.planning.script.builders.worker_pool_builder import build_worker
from nuplan.planning.utils.multithreading.worker_utils import worker_map
def get_traj(args):
output = []
for arg_dict in tqdm(args, desc="Running evaluation"):
input_loader = arg_dict['input_loader']
token = arg_dict['token']
agent = arg_dict['agent']

try:
agent_input = input_loader.get_agent_input_from_token(token)
trajectory = agent.compute_trajectory(agent_input)
output.append([token, trajectory])

except Exception:
logger.warning(f"----------- Agent failed for token {token}:")
traceback.print_exc()

return output

def run_test_evaluation(
cfg: DictConfig,
agent: AbstractAgent,
Expand All @@ -44,6 +64,7 @@ def run_test_evaluation(
Thus, agent.requires_scene has to be False for the agent that is to be evaluated.
"""
)

logger.info("Building Agent Input Loader")
input_loader = SceneLoader(
data_path=data_path,
Expand All @@ -55,30 +76,24 @@ def run_test_evaluation(
)
agent.initialize()

# first stage output
worker = build_worker(cfg)
data_points = [{"input_loader": input_loader, "agent": agent, "token": token,} for token in input_loader.tokens_stage_one]
first_traj = worker_map(worker, get_traj, data_points)
first_stage_output: Dict[str, Trajectory] = {}
for token in tqdm(input_loader.tokens_stage_one, desc="Running first stage evaluation"):
try:
agent_input = input_loader.get_agent_input_from_token(token)
trajectory = agent.compute_trajectory(agent_input)
first_stage_output.update({token: trajectory})
except Exception:
logger.warning(f"----------- Agent failed for token {token}:")
traceback.print_exc()

# second stage output
for first_dict_list in first_traj:
first_stage_output.update({first_dict_list[0]:first_dict_list[1]})


scene_loader_tokens_stage_two = input_loader.reactive_tokens_stage_two

second_stage_output: Dict[str, Trajectory] = {}
for token in tqdm(scene_loader_tokens_stage_two, desc="Running second stage evaluation"):
try:
agent_input = input_loader.get_agent_input_from_token(token)
trajectory = agent.compute_trajectory(agent_input)
second_stage_output.update({token: trajectory})
except Exception:
logger.warning(f"----------- Agent failed for token {token}:")
traceback.print_exc()

data_points = [{"input_loader": input_loader, "agent": agent, "token": token,} for token in scene_loader_tokens_stage_two]
sec_traj = worker_map(worker, get_traj, data_points)

for sec_dict_list in sec_traj:
second_stage_output.update({sec_dict_list[0]:sec_dict_list[1]})

return first_stage_output, second_stage_output

Expand Down
49 changes: 31 additions & 18 deletions navsim/planning/script/run_create_submission_pickle_challenge.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,25 @@
CONFIG_PATH = "config/pdm_scoring"
CONFIG_NAME = "default_run_create_submission_pickle"

from navsim.planning.script.builders.worker_pool_builder import build_worker
from nuplan.planning.utils.multithreading.worker_utils import worker_map
def get_traj(args):
output = []
for arg_dict in tqdm(args, desc="Running evaluation"):
input_loader = arg_dict['input_loader']
token = arg_dict['token']
agent = arg_dict['agent']

try:
agent_input = input_loader.get_agent_input_from_token(token)
trajectory = agent.compute_trajectory(agent_input)
output.append([token, trajectory])

except Exception:
logger.warning(f"----------- Agent failed for token {token}:")
traceback.print_exc()

return output

def run_test_evaluation(
cfg: DictConfig,
Expand Down Expand Up @@ -56,32 +75,26 @@ def run_test_evaluation(
agent.initialize()

# first stage output
worker = build_worker(cfg)
data_points = [{"input_loader": input_loader, "agent": agent, "token": token,} for token in input_loader.tokens_stage_one]
first_traj = worker_map(worker, get_traj, data_points)
first_stage_output: Dict[str, Trajectory] = {}
for token in tqdm(input_loader.tokens_stage_one, desc="Running first stage evaluation"):
try:
agent_input = input_loader.get_agent_input_from_token(token)
trajectory = agent.compute_trajectory(agent_input)
first_stage_output.update({token: trajectory})
except Exception:
logger.warning(f"----------- Agent failed for token {token}:")
traceback.print_exc()

# second stage output
for first_dict_list in first_traj:
first_stage_output.update({first_dict_list[0]:first_dict_list[1]})

# second stage output
scene_loader_tokens_stage_two = input_loader.reactive_tokens_stage_two

second_stage_output: Dict[str, Trajectory] = {}
for token in tqdm(scene_loader_tokens_stage_two, desc="Running second stage evaluation"):
try:
agent_input = input_loader.get_agent_input_from_token(token)
trajectory = agent.compute_trajectory(agent_input)
second_stage_output.update({token: trajectory})
except Exception:
logger.warning(f"----------- Agent failed for token {token}:")
traceback.print_exc()

return first_stage_output, second_stage_output
data_points = [{"input_loader": input_loader, "agent": agent, "token": token,} for token in scene_loader_tokens_stage_two]
sec_traj = worker_map(worker, get_traj, data_points)

for sec_dict_list in sec_traj:
second_stage_output.update({sec_dict_list[0]:sec_dict_list[1]})

return first_stage_output, second_stage_output

@hydra.main(config_path=CONFIG_PATH, config_name=CONFIG_NAME, version_base=None)
def main(cfg: DictConfig) -> None:
Expand Down