def test_scenario_variations_of_social_agents(scenario_root):
iterator = Scenario.variations_for_all_scenario_roots(
[str(scenario_root)], [AGENT_ID]
)
scenarios = list(iterator)
assert len(scenarios) == 6, "3 social agents x 2 missions each "
for s in scenarios:
assert len(s.social_agents) == 4, "4 social agents"
assert len(s.missions) == 5, "4 missions for social agents + 1 for ego"
# Ensure correct social agents are being spawned
all_social_agent_ids = set()
for s in scenarios:
all_social_agent_ids |= set(s.social_agents.keys())
groups = ["group-1", "group-2", "group-3", "group-4"]
speeds = [10, 30, 80]
expected_social_agent_ids = {
SocialAgentId.new(f"non-interactive-agent-{speed}-v0", group=group)
for group, speed in itertools.product(groups, speeds)
}
assert (
len(all_social_agent_ids - expected_social_agent_ids) == 0
), "All the correct social agent IDs were used"
def _airlock_social_vehicle_with_social_agent(
self,
vehicle_id: str,
social_agent_actor: SocialAgentActor,
) -> str:
"""When airlocked. The social agent will receive observations and execute
its policy, however it won't actually operate the vehicle's controller.
"""
self._log.debug(
f"Airlocked vehicle={vehicle_id} with actor={social_agent_actor}")
sim = self._sim
agent_id = BubbleManager._make_social_agent_id(vehicle_id,
social_agent_actor)
if agent_id in sim.agent_manager.social_agent_ids:
# E.g. if agent is a boid and was being re-used
interface = sim.agent_manager.agent_interface_for_agent_id(
agent_id)
else:
social_agent = make_social_agent(
locator=social_agent_actor.agent_locator,
**social_agent_actor.policy_kwargs,
)
interface = social_agent.interface
mission_planner = MissionPlanner(sim.scenario.waypoints,
sim.scenario.road_network)
is_boid = isinstance(social_agent_actor, BoidAgentActor)
vehicle = sim.vehicle_index.prepare_for_agent_control(sim,
vehicle_id,
agent_id,
interface,
mission_planner,
boid=is_boid)
# Setup mission (also used for observations)
route = sim.traffic_sim.vehicle_route(vehicle_id=vehicle.id)
mission = Mission(
start=Start(vehicle.position[:2], vehicle.heading),
goal=PositionalGoal.fromedge(route[-1], sim.scenario.road_network),
)
mission_planner.plan(mission=mission)
if agent_id not in sim.agent_manager.social_agent_ids:
social_agent_data_model = SocialAgent(
id=SocialAgentId.new(social_agent_actor.name),
name=social_agent_actor.name,
mission=mission,
agent_locator=social_agent_actor.agent_locator,
policy_kwargs=social_agent_actor.policy_kwargs,
initial_speed=social_agent_actor.initial_speed,
)
sim.agent_manager.start_social_agent(agent_id, social_agent,
social_agent_data_model)
return agent_id
def _start_social_agent(self, sim, agent_id, social_agent,
social_agent_actor, bubble):
social_agent_data_model = SocialAgent(
id=SocialAgentId.new(social_agent_actor.name),
name=social_agent_actor.name,
is_boid=bubble.is_boid,
is_boid_keep_alive=bubble.keep_alive,
agent_locator=social_agent_actor.agent_locator,
policy_kwargs=social_agent_actor.policy_kwargs,
initial_speed=social_agent_actor.initial_speed,
)
sim.agent_manager.start_social_agent(agent_id, social_agent,
social_agent_data_model)
def start_keep_alive_boid_agents(self, sim):
for bubble in filter(
lambda b: b.is_boid and b.keep_alive, sim.scenario.bubbles
):
actor = bubble.actor
agent_id = BubbleManager._make_boid_social_agent_id(actor)
social_agent = make_social_agent(
locator=actor.agent_locator, **actor.policy_kwargs,
)
actor = bubble.actor
social_agent_data_model = SocialAgent(
id=SocialAgentId.new(actor.name),
name=actor.name,
is_boid=True,
is_boid_keep_alive=True,
agent_locator=actor.agent_locator,
policy_kwargs=actor.policy_kwargs,
initial_speed=actor.initial_speed,
)
self.start_social_agent(agent_id, social_agent, social_agent_data_model)
def to_actor_id(actor, mission_group):
return SocialAgentId.new(actor.name, group=mission_group)
def _discover_social_agents_info(
scenario, ) -> Sequence[Dict[str, SocialAgent]]:
"""Loops through the social agent mission pickles, instantiating corresponding
implementations for the given types. The output is a list of
{agent_id: (mission, locator)}, where each dictionary corresponds to the
social agents to run for a given concrete Scenario (which translates to
"per episode" when swapping).
"""
scenario_root = (scenario.root_filepath if isinstance(
scenario, Scenario) else scenario)
road_map, _ = Scenario.build_map(scenario_root)
social_agents_path = os.path.join(scenario_root, "social_agents")
if not os.path.exists(social_agents_path):
return []
# [ ( missions_file, agent_actor, Mission ) ]
agent_bucketer = []
# like dict.setdefault
def setdefault(l: Sequence[Any], index: int, default):
while len(l) < index + 1:
l.append([])
return l[index]
file_match = os.path.join(social_agents_path, "*.pkl")
for missions_file_path in glob.glob(file_match):
with open(missions_file_path, "rb") as missions_file:
count = 0
missions = pickle.load(missions_file)
for mission_and_actor in missions:
# Each pickle file will contain a list of actor/mission pairs. The pairs
# will most likely be generated in an M:N fashion
# (i.e. A1: M1, A1: M2, A2: M1, A2: M2). The desired behavior is to have
# a single pair per concrete Scenario (which would translate to
# "per episode" when swapping)
assert isinstance(mission_and_actor.actor,
sstudio_types.SocialAgentActor)
actor = mission_and_actor.actor
extracted_mission = Scenario._extract_mission(
mission_and_actor.mission, road_map)
namespace = os.path.basename(missions_file_path)
namespace = os.path.splitext(namespace)[0]
setdefault(agent_bucketer, count, []).append((
SocialAgent(
id=SocialAgentId.new(actor.name, group=namespace),
name=actor.name,
is_boid=False,
is_boid_keep_alive=False,
agent_locator=actor.agent_locator,
policy_kwargs=actor.policy_kwargs,
initial_speed=actor.initial_speed,
),
extracted_mission,
))
count += 1
social_agents_info = []
for l in agent_bucketer:
social_agents_info.append(
{agent.id: (agent, mission)
for agent, mission in l})
return social_agents_info
def to_actor_id(actor, mission_group):
"""Mashes the actor id and mission group to create what needs to be a unique id."""
return SocialAgentId.new(actor.name, group=mission_group)