torchdrivesim.behavior.iai
==========================

.. py:module:: torchdrivesim.behavior.iai


Classes
-------

.. autoapisummary::

   torchdrivesim.behavior.iai.IAIWrapper


Functions
---------

.. autoapisummary::

   torchdrivesim.behavior.iai.unpack_attributes
   torchdrivesim.behavior.iai.iai_initialize
   torchdrivesim.behavior.iai.iai_drive


Module Contents
---------------

.. py:function:: unpack_attributes(attributes) -> torch.Tensor

.. py:function:: iai_initialize(location, agent_count, center=(0, 0), traffic_light_state_history: Optional[List[Dict[int, invertedai.common.TrafficLightState]]] = None) -> (torch.Tensor, torch.Tensor, List[invertedai.common.RecurrentState])

.. py:function:: iai_drive(location: str, agent_states: torch.Tensor, agent_attributes: torch.Tensor, recurrent_states: List[invertedai.common.RecurrentState], traffic_lights_states: Optional[Dict[int, invertedai.common.TrafficLightState]] = None) -> (torch.Tensor, List[invertedai.common.RecurrentState])

.. py:class:: IAIWrapper(simulator: torchdrivesim.simulator.SimulatorInterface, npc_mask: torchdrivesim.simulator.TensorPerAgentType, recurrent_states: List[List], locations: List[str], rear_axis_offset: Optional[torch.Tensor] = None, replay_states: Optional[torch.Tensor] = None, replay_mask: Optional[torch.Tensor] = None, traffic_light_controller: Optional[torchdrivesim.traffic_lights.TrafficLightController] = None, traffic_light_ids: Optional[List[int]] = None)

   Bases: :py:obj:`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.

   :param simulator: existing simulator to wrap
   :param 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.
   :param recurrent_states: A functor of BxAxN tensors with NPC recurrent states, usually obtained from `iai_initialize`.
   :param 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.


   .. py:method:: _get_npc_predictions()


   .. py:method:: _npc_teleport_to()

      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



   .. py:method:: step(action)

      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.



   .. py:method:: to(device) -> typing_extensions.Self

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



   .. py:method:: copy()

      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.



   .. py:method:: select_batch_elements(idx, in_place=True)

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