torchdrivesim.behavior.iai#

Module Contents#

Classes#

IAIWrapper

Uses IAI API to control NPCs, making a call every time step.

Functions#

iai_initialize(location, agent_count[, center])

iai_drive(location, agent_states, agent_attributes, ...)

torchdrivesim.behavior.iai.iai_initialize(location, agent_count, center=(0, 0))[source]#
torchdrivesim.behavior.iai.iai_drive(location: str, agent_states: torch.Tensor, agent_attributes: torch.Tensor, recurrent_states: List)[source]#
class torchdrivesim.behavior.iai.IAIWrapper(simulator: torchdrivesim.simulator.SimulatorInterface, npc_mask: torchdrivesim.simulator.TensorPerAgentType, recurrent_states: List[List], locations: List[str], rear_axis_offset: torch.Tensor | None = None)[source]#

Bases: torchdrivesim.simulator.NPCWrapper

Uses IAI API to control NPCs, making a call every time step. Requires IAI_API_KEY to be set. This wrapper should be applied before any other agent-hiding wrappers.

Parameters:

  • simulator – existing simulator to wrap

  • npc_mask – A functor of tensors with a single dimension of size A, indicating which agents to replay. The tensors can not have batch dimensions.

  • recurrent_states – A functor of BxAxN tensors with NPC recurrent states, usually obtained from iai_initialize.

  • rear_axis_offset – A BxAx1 tensor, concatenated across agent types, specifying the rear axis offset parameter used in the kinematic bicycle model. By default, a realistic value based on the vehicle length is used.

_get_npc_predictions()[source]#
_npc_teleport_to()[source]#

Provides the states to which the NPCs should be set after step, with arbitrary padding for the remaining agents. By default, no teleportation is performed, but subclasses may use it instead of, or on top of defining the NPC action.

Returns:

a functor of BxAxSt tensors, where A is the number of agents in the inner simulator, or None if no teleportation is required

step(action)[source]#

Runs the simulation for one step with given agent actions. Input is a functor of BxAxAc tensors, where Ac is determined by the kinematic model.

to(device) → typing_extensions.Self[source]#

Modifies the simulator in-place, putting all tensors on the device provided.

copy()[source]#

Duplicates this simulator, allowing for independent subsequent execution. The copy is relatively shallow, in that the tensors are the same objects but dictionaries referring to them are shallowly copied.

select_batch_elements(idx, in_place=True)[source]#

Picks selected elements of the batch. The input is a tensor of indices into the batch dimension.