torchdrivesim.simulator

Module Contents

Classes

CollisionMetric	Method used to calculate collisions between agents.
TorchDriveConfig	Top-level configuration for a TorchDriveSim simulator.
AgentTypeFunctor	Lifts functions operating on individual agent type to functions operating across all agent types.
SingletonAgentTypeFunctor	Trivial functor where only one agent type is used.
DictAgentTypeFunctor	Arguments are packaged as dictionaries, with keys being agent type names.
SimulatorInterface	Abstract interface for a 2D differentiable driving simulator.
Simulator	Base simulator, where the agent functor is a dictionary indexed with agent type.
SimulatorWrapper	Modifies the behavior of an existing simulator, itself acting like a simulator.
NPCWrapper	Designates a certain subset of agents as non-playable characters (NPCs) and removes them from
HomogeneousWrapper	Removes distinction between agent types, replacing collections with single tensors
RecordingWrapper	Doesn't modify the behavior of the simulator, but records some information at each time step.
BirdviewRecordingWrapper	Records a visualization of the simulation state at each step.
TrajectoryVisualizationWrapper	Records birdviews with additional trajectory visualizations.
SelectiveWrapper	Only exposes a subset of agents from the inner simulator.
BoundedRegionWrapper	A variant of SelectiveWrapper that exposes agents contained within a given polygon defining the area of interest.
NoReentryBoundedRegionWrapper	A variant of BoundedRegionWrapper that does not allow reentry.

Attributes

logger
TensorPerAgentType
IntPerAgentType

torchdrivesim.simulator.logger

class torchdrivesim.simulator.CollisionMetric

Bases: enum.Enum

Method used to calculate collisions between agents.

iou = 'iou'

discs = 'discs'

nograd = 'nograd'

nograd_pytorch3d = 'nograd-pytorch3d'

class torchdrivesim.simulator.TorchDriveConfig

Top-level configuration for a TorchDriveSim simulator.

renderer: torchdrivesim.rendering.RendererConfig

single_agent_rendering: bool = False

collision_metric: CollisionMetric

offroad_threshold: float = 0.5

left_handed_coordinates: bool = False

torchdrivesim.simulator.TensorPerAgentType

torchdrivesim.simulator.IntPerAgentType

class torchdrivesim.simulator.AgentTypeFunctor

Bases: abc.ABC

Lifts functions operating on individual agent type to functions operating across all agent types. Children of this class specify how the collection of agents is represented.

abstract fmap(f, *args)

Applies a given function of any number of arguments to all agent types.

to_device(tensor: TensorPerAgentType, device: torch.device) → TensorPerAgentType

Applies .to to each element.

__call__(f, *args)

class torchdrivesim.simulator.SingletonAgentTypeFunctor

Bases: AgentTypeFunctor

Trivial functor where only one agent type is used. No packaging is used for the arguments.

fmap(f, *args)

Applies a given function of any number of arguments to all agent types.

class torchdrivesim.simulator.DictAgentTypeFunctor(agent_types: List[str])

Bases: AgentTypeFunctor

Arguments are packaged as dictionaries, with keys being agent type names.

fmap(f, *args)

Applies a given function of any number of arguments to all agent types.

class torchdrivesim.simulator.SimulatorInterface

Abstract interface for a 2D differentiable driving simulator.

abstract property agent_functor: AgentTypeFunctor

Defines how to apply functions across agent types.

abstract property agent_types: List[str] | None

List of agent types used by this simulator, or None if only one type used.

abstract property action_size: IntPerAgentType

Defines the size of the action space for each agent type.

abstract property batch_size: int

property agent_count: IntPerAgentType

How many agents of each type are there in the simulation. This counts the available slots, not taking present masks into consideration.

across_agent_types(f, *args)

Applies a given per-agent-type operation across all agent types. Subsequent arguments should be given in the format of multiple agent collection that this class uses, which is typically a dictionary with keys being agent type names.

abstract to(device) → typing_extensions.Self

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

abstract copy() → typing_extensions.Self

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.

abstract extend(n: int, in_place: bool = True) → typing_extensions.Self

Multiplies the first batch dimension by the given number. Like in pytorch3d, this is equivalent to introducing extra batch dimension on the right and then flattening.

abstract select_batch_elements(idx, in_place=True) → typing_extensions.Self

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

__getitem__(item: torch.Tensor) → typing_extensions.Self

Allows indexing syntax. item should be an iterable collection of indices.

abstract get_world_center() → torch.Tensor

Returns a Bx2 tensor with the coordinates of the map center.

abstract get_state() → TensorPerAgentType

Returns a functor of BxAxSt tensors representing current agent states.

abstract get_agent_size() → TensorPerAgentType

Returns a functor of BxAx2 tensors representing agent length and width.

abstract get_agent_type() → TensorPerAgentType

Returns a functor of BxA long tensors containing agent type indexes relative to the list containing all agent types

as returned by SimulatorInterface.agent_types.

abstract get_present_mask() → TensorPerAgentType

Returns a functor of BxA boolean tensors indicating which agents are currently present in the simulation.

abstract get_all_agents_absolute() → TensorPerAgentType

Returns a functor of BxAx6 tensors, where the last dimension contains the following information: x, y, psi, length, width, present. Typically used to implement non-visual observation modalities.

abstract get_all_agents_relative(exclude_self: bool = True) → TensorPerAgentType

Returns a functor of BxAxAllx6 tensors, specifying for each of A agents the relative position about the other agents. ‘All’ is the number of all agents in the simulation, including hidden ones, across all agent types. If exclude_self is set, for each agent in A, that agent itself is removed from All. The final dimension has the same meaning as in get_all_agents_absolute, except now the positions and orientations are relative to the specified agent.

get_traffic_controls() → Dict[str, torchdrivesim.traffic_controls.BaseTrafficControl]

Produces all traffic controls existing in the simulation, grouped by type.

abstract get_innermost_simulator() → typing_extensions.Self

Returns the innermost concrete Simulator object. The type signature is misleading due to Python limitations.

abstract get_waypoints() → TensorPerAgentType

Returns a functor of BxAxMx2 tensors representing current agent waypoints.

abstract get_waypoints_mask() → TensorPerAgentType

Returns a functor of BxAxM boolean tensors representing current agent waypoints present mask.

abstract step(agent_action: TensorPerAgentType) → None

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.

abstract set_state(agent_state: TensorPerAgentType, mask: TensorPerAgentType | None = None) → None

Arbitrarily set the state of the agents, without advancing the simulation. The change is effective immediately, without waiting for the next step.

Parameters:

• agent_state – a functor of BxAx4 tensors with agent states

• mask – a functor of BxA boolean tensors, deciding which agent states to update; all by default

abstract update_present_mask(present_mask: TensorPerAgentType) → None

Sets the present mask of agents to the provided value.

Parameters:

present_mask – a functor of BxA boolean tensors

abstract fit_action(future_state: TensorPerAgentType, current_state: TensorPerAgentType | None = None) → TensorPerAgentType

Computes an action that would (aproximately) produce the desired state.

Parameters:

• future_state – a functor of BxAx4 tensors defining the desired state

• current_state – if different from the current simulation state, in the same format as future state

Returns:

a functor of BxAxAc tensors

abstract render(camera_xy: torch.Tensor, camera_psi: torch.Tensor, res: torchdrivesim.utils.Resolution | None = None, rendering_mask: TensorPerAgentType | None = None, fov: float | None = None, waypoints: torch.Tensor | None = None, waypoints_rendering_mask: torch.Tensor | None = None) → torch.Tensor

Renders the world from bird’s eye view using cameras in given positions.

Parameters:

• camera_xy – BxNx2 tensor of x-y positions for N cameras

• camera_psi – BxNx1 tensor of orientations for N cameras

• res – desired image resolution (only square resolutions are supported; by default use value from config)

• rendering_mask – functor of BxNxA tensors, indicating which agents should be rendered each camera

• fov – the field of view of the resulting image in meters (by default use value from config)

• waypoints – BxNxMx2 tensor of M waypoints per camera (x,y)

• waypoints_rendering_mask – BxNxM tensor of M waypoint masks per camera, indicating which waypoints should be rendered

Returns:

BxNxCxHxW tensor of resulting RGB images for each camera

render_egocentric(ego_rotate: bool = True, res: torchdrivesim.utils.Resolution | None = None, fov: float | None = None) → TensorPerAgentType

Renders the world using cameras placed on each agent.

Parameters:

• ego_rotate – whether to orient the cameras such that the ego agent faces up in the image

• res – desired image resolution (only square resolutions are supported; by default use value from config)

• fov – the field of view of the resulting image in meters (by default use value from config)

Returns:

a functor of BxAxCxHxW tensors of resulting RGB images for each agent.

abstract compute_offroad() → TensorPerAgentType

Offroad metric for each agent, defined as the distance to the road mesh. See torchdrivesim.infractions.offroad_infraction_loss for details.

Returns:

a functor of BxA tensors

abstract compute_wrong_way() → TensorPerAgentType

Wrong-way metric for each agent, based on the inner product between the agent and lane direction. See torchdrivesim.infractions.lanelet_orientation_loss for details.

Returns:

a functor of BxA tensors

compute_traffic_lights_violations() → TensorPerAgentType

Boolean value indicating whether each agent is committing a traffic light violation. See torchdrivesim.infractions.traffic_controls.TrafficLightControl.compute_violations for details.

Returns:

a functor of BxA tensors

abstract _compute_collision_of_single_agent(box: torch.Tensor, remove_self_overlap: torch.Tensor | None = None, agent_types: List[str] | None = None) → torch.Tensor

Computes the collision metric for an agent specified as a bounding box. Includes collisions with all agents in the simulation, including the ones not exposed through the interface of this class. Used with discs and iou metrics.

Parameters:

• box – Bx5 tensor, with the last dimension being (x,y,length,width,psi).

• remove_self_overlap – B boolean tensor, where if the input agent is present in the simulation, set this to subtract self-overlap. By default it is assumed that self overlapping exists and will be removed.

• agent_types – An optional list of specific agent types for computing collisions with. By default all available agent types will be used.

Returns:

a tensor with a single dimension of B elements

abstract _compute_collision_of_multi_agents(mask: torch.Tensor | None = None) → TensorPerAgentType

Computes the collision metric for selected (default all) agents in the simulation. Includes collisions with all agents in the simulation, including the ones not exposed through the interface of this class. Used with nograd and nograd-pytorch3d metrics.

Parameters:

mask – a functor of BxA boolean tensors, indicating for which agents to compute the loss (by default use present mask)

Returns:

a functor of BxA tensors

compute_collision(agent_types: List[str] | None = None) → TensorPerAgentType

Compute the collision metric for agents exposed through the interface of this class. Includes collisions with agents not exposed through the interface. Collisions are defined as overlap of agents’ bounding boxes, with details determined by the specific method chosen in the config.

Parameters:

agent_types – An optional list of specific agent types for computing collisions with. Not supported by the collision metrics nograd and nograd-pytorch3d.

Returns:

a functor of BxA tensors

class torchdrivesim.simulator.Simulator(road_mesh: torchdrivesim.mesh.BirdviewMesh, kinematic_model: Dict[str, torchdrivesim.kinematic.KinematicModel], agent_size: Dict[str, torch.Tensor], initial_present_mask: Dict[str, torch.Tensor], cfg: TorchDriveConfig, renderer: torchdrivesim.rendering.BirdviewRenderer | None = None, lanelet_map: List[torchdrivesim.lanelet2.LaneletMap | None] | None = None, recenter_offset: torch.Tensor | None = None, internal_time: int = 0, traffic_controls: Dict[str, torchdrivesim.traffic_controls.BaseTrafficControl] | None = None, waypoint_goals: torchdrivesim.goals.WaypointGoal | None = None)

Bases: SimulatorInterface

Base simulator, where the agent functor is a dictionary indexed with agent type.

Parameters:

• road_mesh – a mesh indicating the driveable area

• kinematic_model – determines the action space, constraints, and the initial state of all agents

• agent_size – a functor of Bx2 tensors indicating agent length and width

• initial_present_mask – a functor of BxA tensors indicating which agents are initially present and not padding

• cfg – holds various configuration options

• renderer – specify if using a non-standard renderer or static meshes beyond the road mesh (default from config)

• lanelet_map – provide the map to compute orientation losses, one map per batch element where available

• recenter_offset – if the coordinate system from lanelet_map was shifted, this value will be used to shift it back

• internal_time – initial value for step counter

• traffic_controls – applicable traffic controls by type

property agent_functor: DictAgentTypeFunctor

Defines how to apply functions across agent types.

property agent_types

List of agent types used by this simulator, or None if only one type used.

property action_size: IntPerAgentType

Defines the size of the action space for each agent type.

property batch_size: int

to(device)

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

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.

extend(n, in_place=True)

Multiplies the first batch dimension by the given number. Like in pytorch3d, this is equivalent to introducing extra batch dimension on the right and then flattening.

select_batch_elements(idx, in_place=True)

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

validate_agent_types()

validate_tensor_shapes()

validate_agent_count(count_dict)

get_world_center()

Returns a Bx2 tensor with the coordinates of the map center.

get_state()

Returns a functor of BxAxSt tensors representing current agent states.

get_waypoints()

Returns a functor of BxAxMx2 tensors representing current agent waypoints.

get_waypoints_mask()

Returns a functor of BxAxM boolean tensors representing current agent waypoints present mask.

compute_wrong_way()

Wrong-way metric for each agent, based on the inner product between the agent and lane direction. See torchdrivesim.infractions.lanelet_orientation_loss for details.

Returns:

a functor of BxA tensors

get_agent_size()

Returns a functor of BxAx2 tensors representing agent length and width.

get_agent_type()

Returns a functor of BxA long tensors containing agent type indexes relative to the list containing all agent types

as returned by SimulatorInterface.agent_types.

get_present_mask()

Returns a functor of BxA boolean tensors indicating which agents are currently present in the simulation.

get_all_agents_absolute()

Returns a functor of BxAx6 tensors, where the last dimension contains the following information: x, y, psi, length, width, present. Typically used to implement non-visual observation modalities.

get_all_agents_relative(exclude_self=True)

Returns a functor of BxAxAllx6 tensors, specifying for each of A agents the relative position about the other agents. ‘All’ is the number of all agents in the simulation, including hidden ones, across all agent types. If exclude_self is set, for each agent in A, that agent itself is removed from All. The final dimension has the same meaning as in get_all_agents_absolute, except now the positions and orientations are relative to the specified agent.

get_innermost_simulator() → typing_extensions.Self

Returns the innermost concrete Simulator object. The type signature is misleading due to Python limitations.

step(agent_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.

set_state(agent_state, mask=None)

Arbitrarily set the state of the agents, without advancing the simulation. The change is effective immediately, without waiting for the next step.

Parameters:

• agent_state – a functor of BxAx4 tensors with agent states

• mask – a functor of BxA boolean tensors, deciding which agent states to update; all by default

update_present_mask(present_mask)

Sets the present mask of agents to the provided value.

Parameters:

present_mask – a functor of BxA boolean tensors

fit_action(future_state, current_state=None)

Computes an action that would (aproximately) produce the desired state.

Parameters:

• future_state – a functor of BxAx4 tensors defining the desired state

• current_state – if different from the current simulation state, in the same format as future state

Returns:

a functor of BxAxAc tensors

render(camera_xy, camera_psi, res=None, rendering_mask=None, fov=None, waypoints=None, waypoints_rendering_mask=None)

Renders the world from bird’s eye view using cameras in given positions.

Parameters:

• camera_xy – BxNx2 tensor of x-y positions for N cameras

• camera_psi – BxNx1 tensor of orientations for N cameras

• res – desired image resolution (only square resolutions are supported; by default use value from config)

• rendering_mask – functor of BxNxA tensors, indicating which agents should be rendered each camera

• fov – the field of view of the resulting image in meters (by default use value from config)

• waypoints – BxNxMx2 tensor of M waypoints per camera (x,y)

• waypoints_rendering_mask – BxNxM tensor of M waypoint masks per camera, indicating which waypoints should be rendered

Returns:

BxNxCxHxW tensor of resulting RGB images for each camera

compute_offroad()

Offroad metric for each agent, defined as the distance to the road mesh. See torchdrivesim.infractions.offroad_infraction_loss for details.

Returns:

a functor of BxA tensors

_compute_collision_of_single_agent(box, remove_self_overlap=None, agent_types=None)

Computes the collision metric for an agent specified as a bounding box. Includes collisions with all agents in the simulation, including the ones not exposed through the interface of this class. Used with discs and iou metrics.

Parameters:

• box – Bx5 tensor, with the last dimension being (x,y,length,width,psi).

• remove_self_overlap – B boolean tensor, where if the input agent is present in the simulation, set this to subtract self-overlap. By default it is assumed that self overlapping exists and will be removed.

• agent_types – An optional list of specific agent types for computing collisions with. By default all available agent types will be used.

Returns:

a tensor with a single dimension of B elements

_compute_collision_of_multi_agents(mask=None)

Computes the collision metric for selected (default all) agents in the simulation. Includes collisions with all agents in the simulation, including the ones not exposed through the interface of this class. Used with nograd and nograd-pytorch3d metrics.

Parameters:

mask – a functor of BxA boolean tensors, indicating for which agents to compute the loss (by default use present mask)

Returns:

a functor of BxA tensors

class torchdrivesim.simulator.SimulatorWrapper(simulator: SimulatorInterface)

Bases: SimulatorInterface

Modifies the behavior of an existing simulator, itself acting like a simulator. This base class simply delegates all method calls to the inner simulator.

property agent_functor: AgentTypeFunctor

Defines how to apply functions across agent types.

property agent_types

List of agent types used by this simulator, or None if only one type used.

property action_size: IntPerAgentType

Defines the size of the action space for each agent type.

property batch_size: int

to(device) → typing_extensions.Self

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

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.

extend(n, in_place=True)

Multiplies the first batch dimension by the given number. Like in pytorch3d, this is equivalent to introducing extra batch dimension on the right and then flattening.

select_batch_elements(idx, in_place=True)

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

get_world_center()

Returns a Bx2 tensor with the coordinates of the map center.

get_state()

Returns a functor of BxAxSt tensors representing current agent states.

get_agent_size()

Returns a functor of BxAx2 tensors representing agent length and width.

get_agent_type()

Returns a functor of BxA long tensors containing agent type indexes relative to the list containing all agent types

as returned by SimulatorInterface.agent_types.

get_present_mask()

Returns a functor of BxA boolean tensors indicating which agents are currently present in the simulation.

get_waypoints()

Returns a functor of BxAxMx2 tensors representing current agent waypoints.

get_waypoints_mask()

Returns a functor of BxAxM boolean tensors representing current agent waypoints present mask.

get_all_agents_absolute()

Returns a functor of BxAx6 tensors, where the last dimension contains the following information: x, y, psi, length, width, present. Typically used to implement non-visual observation modalities.

get_all_agents_relative(exclude_self=True)

Returns a functor of BxAxAllx6 tensors, specifying for each of A agents the relative position about the other agents. ‘All’ is the number of all agents in the simulation, including hidden ones, across all agent types. If exclude_self is set, for each agent in A, that agent itself is removed from All. The final dimension has the same meaning as in get_all_agents_absolute, except now the positions and orientations are relative to the specified agent.

get_innermost_simulator() → Simulator

Returns the innermost concrete Simulator object. The type signature is misleading due to Python limitations.

step(*args, **kwargs)

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.

set_state(*args, **kwargs)

Arbitrarily set the state of the agents, without advancing the simulation. The change is effective immediately, without waiting for the next step.

Parameters:

• agent_state – a functor of BxAx4 tensors with agent states

• mask – a functor of BxA boolean tensors, deciding which agent states to update; all by default

update_present_mask(*args, **kwargs)

Sets the present mask of agents to the provided value.

Parameters:

present_mask – a functor of BxA boolean tensors

fit_action(*args, **kwargs)

Computes an action that would (aproximately) produce the desired state.

Parameters:

• future_state – a functor of BxAx4 tensors defining the desired state

• current_state – if different from the current simulation state, in the same format as future state

Returns:

a functor of BxAxAc tensors

compute_offroad()

Offroad metric for each agent, defined as the distance to the road mesh. See torchdrivesim.infractions.offroad_infraction_loss for details.

Returns:

a functor of BxA tensors

_compute_collision_of_single_agent(box, remove_self_overlap=None, agent_types=None)

Computes the collision metric for an agent specified as a bounding box. Includes collisions with all agents in the simulation, including the ones not exposed through the interface of this class. Used with discs and iou metrics.

Parameters:

• box – Bx5 tensor, with the last dimension being (x,y,length,width,psi).

• remove_self_overlap – B boolean tensor, where if the input agent is present in the simulation, set this to subtract self-overlap. By default it is assumed that self overlapping exists and will be removed.

• agent_types – An optional list of specific agent types for computing collisions with. By default all available agent types will be used.

Returns:

a tensor with a single dimension of B elements

_compute_collision_of_multi_agents(mask=None)

Computes the collision metric for selected (default all) agents in the simulation. Includes collisions with all agents in the simulation, including the ones not exposed through the interface of this class. Used with nograd and nograd-pytorch3d metrics.

Parameters:

mask – a functor of BxA boolean tensors, indicating for which agents to compute the loss (by default use present mask)

Returns:

a functor of BxA tensors

compute_wrong_way()

Wrong-way metric for each agent, based on the inner product between the agent and lane direction. See torchdrivesim.infractions.lanelet_orientation_loss for details.

Returns:

a functor of BxA tensors

render(camera_xy, camera_psi, res=None, rendering_mask=None, fov=None, waypoints=None, waypoints_rendering_mask=None)

Renders the world from bird’s eye view using cameras in given positions.

Parameters:

• camera_xy – BxNx2 tensor of x-y positions for N cameras

• camera_psi – BxNx1 tensor of orientations for N cameras

• res – desired image resolution (only square resolutions are supported; by default use value from config)

• rendering_mask – functor of BxNxA tensors, indicating which agents should be rendered each camera

• fov – the field of view of the resulting image in meters (by default use value from config)

• waypoints – BxNxMx2 tensor of M waypoints per camera (x,y)

• waypoints_rendering_mask – BxNxM tensor of M waypoint masks per camera, indicating which waypoints should be rendered

Returns:

BxNxCxHxW tensor of resulting RGB images for each camera

class torchdrivesim.simulator.NPCWrapper(simulator: SimulatorInterface, npc_mask: TensorPerAgentType)

Bases: SimulatorWrapper

Designates a certain subset of agents as non-playable characters (NPCs) and removes them from the simulator interface, although they remain in simulation and can be interacted with. Note that the designation of which agents are NPCs needs to be the same across batch elements. Subclasses implement specific policies controlling NPC behaviors. At a minimum, they should implement either _get_npc_action or _npc_teleport_to.

Parameters:

npc_mask – A functor of tensors with a single dimension of size A, indicating which agents to replay.

_update_npc_present_mask() → TensorPerAgentType

Computes updated present masks for NPCs, with arbitrary padding for the remaining agents. By default, leaves present masks unchanged.

Returns:

a functor of BxA boolean tensors, where A is the number of agents in the inner simulator

_get_npc_action() → TensorPerAgentType

Computes the actions for NPCs, with arbitrary padding for actions of the remaining agents. By default, the actions are all zeros, but subclasses can implement more intelligent behavior.

Returns:

a functor of BxAxAc tensors, where A is the number of agents in the inner simulator

_npc_teleport_to() → TensorPerAgentType | None

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

to(device) → typing_extensions.Self

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

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.

validate_agent_count(count_dict)

get_state()

Returns a functor of BxAxSt tensors representing current agent states.

get_waypoints()

Returns a functor of BxAxMx2 tensors representing current agent waypoints.

get_waypoints_mask()

Returns a functor of BxAxM boolean tensors representing current agent waypoints present mask.

get_agent_size()

Returns a functor of BxAx2 tensors representing agent length and width.

get_agent_type()

Returns a functor of BxA long tensors containing agent type indexes relative to the list containing all agent types

as returned by SimulatorInterface.agent_types.

get_present_mask()

Returns a functor of BxA boolean tensors indicating which agents are currently present in the simulation.

get_all_agents_relative(exclude_self=True)

Returns a functor of BxAxAllx6 tensors, specifying for each of A agents the relative position about the other agents. ‘All’ is the number of all agents in the simulation, including hidden ones, across all agent types. If exclude_self is set, for each agent in A, that agent itself is removed from All. The final dimension has the same meaning as in get_all_agents_absolute, except now the positions and orientations are relative to the specified agent.

set_state(agent_state, mask=None)

Arbitrarily set the state of the agents, without advancing the simulation. The change is effective immediately, without waiting for the next step.

Parameters:

• agent_state – a functor of BxAx4 tensors with agent states

• mask – a functor of BxA boolean tensors, deciding which agent states to update; all by default

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.

update_present_mask(present_mask)

Sets the present mask of agents to the provided value.

Parameters:

present_mask – a functor of BxA boolean tensors

render(camera_xy, camera_psi, res=None, rendering_mask=None, fov=None, waypoints=None, waypoints_rendering_mask=None)

Renders the world from bird’s eye view using cameras in given positions.

Parameters:

• camera_xy – BxNx2 tensor of x-y positions for N cameras

• camera_psi – BxNx1 tensor of orientations for N cameras

• res – desired image resolution (only square resolutions are supported; by default use value from config)

• rendering_mask – functor of BxNxA tensors, indicating which agents should be rendered each camera

• fov – the field of view of the resulting image in meters (by default use value from config)

• waypoints – BxNxMx2 tensor of M waypoints per camera (x,y)

• waypoints_rendering_mask – BxNxM tensor of M waypoint masks per camera, indicating which waypoints should be rendered

Returns:

BxNxCxHxW tensor of resulting RGB images for each camera

fit_action(future_state, current_state=None)

Computes an action that would (aproximately) produce the desired state.

Parameters:

• future_state – a functor of BxAx4 tensors defining the desired state

• current_state – if different from the current simulation state, in the same format as future state

Returns:

a functor of BxAxAc tensors

compute_offroad()

Offroad metric for each agent, defined as the distance to the road mesh. See torchdrivesim.infractions.offroad_infraction_loss for details.

Returns:

a functor of BxA tensors

compute_wrong_way()

Wrong-way metric for each agent, based on the inner product between the agent and lane direction. See torchdrivesim.infractions.lanelet_orientation_loss for details.

Returns:

a functor of BxA tensors

_compute_collision_of_multi_agents(mask=None)

Computes the collision metric for selected (default all) agents in the simulation. Includes collisions with all agents in the simulation, including the ones not exposed through the interface of this class. Used with nograd and nograd-pytorch3d metrics.

Parameters:

mask – a functor of BxA boolean tensors, indicating for which agents to compute the loss (by default use present mask)

Returns:

a functor of BxA tensors

static extend_tensor(x, batch_dims)

class torchdrivesim.simulator.HomogeneousWrapper(simulator: SimulatorInterface)

Bases: SimulatorWrapper

Removes distinction between agent types, replacing collections with single tensors concatenated across the agent dimension. Only safe to use when all agents share the same kinematic model.

property agent_functor: SingletonAgentTypeFunctor

Defines how to apply functions across agent types.

property agent_types

List of agent types used by this simulator, or None if only one type used.

property action_size: IntPerAgentType

Defines the size of the action space for each agent type.

agent_concat(d, agent_dim, check_agent_count=True)

Concatenates a collection of agent tensors along the agent dimension.

agent_split(x, agent_dim)

Reverse of agent_concat, splits a tensor into a collection of agent tensors.

get_state()

Returns a functor of BxAxSt tensors representing current agent states.

get_agent_size()

Returns a functor of BxAx2 tensors representing agent length and width.

get_agent_type()

Returns a functor of BxA long tensors containing agent type indexes relative to the list containing all agent types

as returned by SimulatorInterface.agent_types.

get_present_mask()

Returns a functor of BxA boolean tensors indicating which agents are currently present in the simulation.

get_waypoints()

Returns a functor of BxAxMx2 tensors representing current agent waypoints.

get_waypoints_mask()

Returns a functor of BxAxM boolean tensors representing current agent waypoints present mask.

get_all_agents_absolute()

Returns a functor of BxAx6 tensors, where the last dimension contains the following information: x, y, psi, length, width, present. Typically used to implement non-visual observation modalities.

get_all_agents_relative(exclude_self=True)

Returns a functor of BxAxAllx6 tensors, specifying for each of A agents the relative position about the other agents. ‘All’ is the number of all agents in the simulation, including hidden ones, across all agent types. If exclude_self is set, for each agent in A, that agent itself is removed from All. The final dimension has the same meaning as in get_all_agents_absolute, except now the positions and orientations are relative to the specified agent.

set_state(agent_state, mask=None)

Arbitrarily set the state of the agents, without advancing the simulation. The change is effective immediately, without waiting for the next step.

Parameters:

• agent_state – a functor of BxAx4 tensors with agent states

• mask – a functor of BxA boolean tensors, deciding which agent states to update; all by default

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.

update_present_mask(present_mask)

Sets the present mask of agents to the provided value.

Parameters:

present_mask – a functor of BxA boolean tensors

fit_action(future_state, current_state=None)

Computes an action that would (aproximately) produce the desired state.

Parameters:

• future_state – a functor of BxAx4 tensors defining the desired state

• current_state – if different from the current simulation state, in the same format as future state

Returns:

a functor of BxAxAc tensors

compute_offroad()

Offroad metric for each agent, defined as the distance to the road mesh. See torchdrivesim.infractions.offroad_infraction_loss for details.

Returns:

a functor of BxA tensors

compute_wrong_way()

Wrong-way metric for each agent, based on the inner product between the agent and lane direction. See torchdrivesim.infractions.lanelet_orientation_loss for details.

Returns:

a functor of BxA tensors

_compute_collision_of_multi_agents(mask=None)

Computes the collision metric for selected (default all) agents in the simulation. Includes collisions with all agents in the simulation, including the ones not exposed through the interface of this class. Used with nograd and nograd-pytorch3d metrics.

Parameters:

mask – a functor of BxA boolean tensors, indicating for which agents to compute the loss (by default use present mask)

Returns:

a functor of BxA tensors

render(camera_xy, camera_psi, res=None, rendering_mask=None, fov=None, waypoints=None, waypoints_rendering_mask=None)

Renders the world from bird’s eye view using cameras in given positions.

Parameters:

• camera_xy – BxNx2 tensor of x-y positions for N cameras

• camera_psi – BxNx1 tensor of orientations for N cameras

• res – desired image resolution (only square resolutions are supported; by default use value from config)

• rendering_mask – functor of BxNxA tensors, indicating which agents should be rendered each camera

• fov – the field of view of the resulting image in meters (by default use value from config)

• waypoints – BxNxMx2 tensor of M waypoints per camera (x,y)

• waypoints_rendering_mask – BxNxM tensor of M waypoint masks per camera, indicating which waypoints should be rendered

Returns:

BxNxCxHxW tensor of resulting RGB images for each camera

class torchdrivesim.simulator.RecordingWrapper(simulator: SimulatorInterface, record_functions: Dict[str, Callable[[SimulatorInterface], Any]], initial_recording: bool = True)

Bases: SimulatorWrapper

Doesn’t modify the behavior of the simulator, but records some information at each time step.

Parameters:

• record_functions – those functions are called on self after each step and their outputs are recorded

• initial_recording – whether to record values when constructing the object, before the first step

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.

record()

Appends the results based on a current state to records.

get_records()

Returns the dictionary of recorded values. Each entry is a list with results from subsequent invocations of record.

step(agent_action, record=True)

By default performs recording at the end of step, but this can be disabled to manually control when to record.

class torchdrivesim.simulator.BirdviewRecordingWrapper(simulator: SimulatorInterface, res: torchdrivesim.utils.Resolution = Resolution(1024, 1024), fov: float = 100, camera_xy: torch.Tensor | None = None, camera_psi: torch.Tensor | None = None, initial_recording: bool = True, to_cpu: bool = False)

Bases: RecordingWrapper

Records a visualization of the simulation state at each step. The visualization is obtained by calling self.render.

Parameters:

• simulator – a simulator object to wrap

• res – resolution for the recorded images

• fov – field of view for the recorded images

• camera_xy – a Bx2 tensor of x-y camera positions (world center by default)

• camera_psi – a Bx1 tensor of camera orientations

• initial_recording – whether to record an image before the first step

• to_cpu – whether to move recorded images to CPU memory, releasing GPU memory

to(device) → typing_extensions.Self

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

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.

extend(n, in_place=True)

Multiplies the first batch dimension by the given number. Like in pytorch3d, this is equivalent to introducing extra batch dimension on the right and then flattening.

select_batch_elements(idx, in_place=True)

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

get_birdviews(stack: bool = False) → torch.Tensor | List[torch.Tensor]

Extracts recorded images.

Parameters:

stack – whether to concatenate results across time, which is placed after batch dimension

Returns:

a BxTxCxHxW tensor of RGB images if stack is set, a list of BxCxHxW images with the same content otherwise

save_gif(filename: str, batch_index: int = 0, fps: float = 10) → None

Saves a GIF to disk using all birdviews for a selected example from the batch.

save_png(filename: str, frame: int, batch_index: int = 0) → None

Saves the given frame as a PNG to disk.

class torchdrivesim.simulator.TrajectoryVisualizationWrapper(simulator: SimulatorInterface, ground_truth: torch.Tensor, predictions: torch.Tensor, gt_present_masks: torch.Tensor, predictions_present_masks: torch.Tensor, trajectory_present_masks: torch.Tensor, res: torchdrivesim.utils.Resolution = Resolution(1024, 1024), fov: float = 100, camera_xy: torch.Tensor | None = None, camera_psi: torch.Tensor | None = None, to_cpu: bool = False)

Bases: BirdviewRecordingWrapper

Records birdviews with additional trajectory visualizations. Those visualizations will be visible to agents relying on rendered observations, so typically this wrapper is used with replay states. The wrapper visualizes one ground truth trajectory and S predictions for each agent. Visualized trajectories should be concatenated across the agent dimension.

Parameters:

• ground_truth – BxAxTxSt tensor with ground truth trajectories for all agents

• predictions – BxSxAxTxSt tensor with predicted trajectories

• gt_present_masks – BxSxAxT boolean tensor, indicating which points in predictions should be visualized

• predictions_present_masks – BxAxT boolean tensor, indicating which points in ground truth should be visualized

class torchdrivesim.simulator.SelectiveWrapper(simulator: SimulatorInterface, exposed_agent_limit: int, default_action: TensorPerAgentType)

Bases: SimulatorWrapper

Only exposes a subset of agents from the inner simulator. The selection may change over time, but there is a fixed number of agent slots available. Which slots are actually used is indicated by the presence mask. The base class exposes the first E agents and the children should override update_exposed_agents and get_present_mask to customize it. Typically, hidden agents are either absent or their state is set externally, but by default the default_action is executed for them.

Parameters:

• simulator – existing simulator to wrap

• exposed_agent_limit – denoted as E in comments, sets the number of slots for exposed agents

• default_action – a functor of BxAxAc tensors, defining the default action to use for non-exposed agents

update_exposed_agents() → None

There should always be E exposed agents, although some of them may be marked as absent.

get_exposed_agents() → TensorPerAgentType

Returns:

a functor of BxE int tensors of indices of exposed agents

is_exposed() → TensorPerAgentType

Returns:

a functor of BxA boolean tensors indicating which agents are exposed

_restrict_tensor(tensor: TensorPerAgentType, agent_dim: int) → TensorPerAgentType

Selects exposed agents from a given tensor, along the supplied agent dimension.

_extend_tensor(tensor: TensorPerAgentType, padding: TensorPerAgentType) → TensorPerAgentType

Given a tensor of exposed agents, constructs a tensor of all agents, using supplied padding. Agent dimension should be the penultimate one.

to(device) → typing_extensions.Self

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

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.

extend(n, in_place=True)

Multiplies the first batch dimension by the given number. Like in pytorch3d, this is equivalent to introducing extra batch dimension on the right and then flattening.

select_batch_elements(idx, in_place=True)

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

get_state()

Returns a functor of BxAxSt tensors representing current agent states.

get_present_mask()

Returns a functor of BxA boolean tensors indicating which agents are currently present in the simulation.

get_agent_size()

Returns a functor of BxAx2 tensors representing agent length and width.

get_agent_type()

Returns a functor of BxA long tensors containing agent type indexes relative to the list containing all agent types

as returned by SimulatorInterface.agent_types.

get_waypoints()

Returns a functor of BxAxMx2 tensors representing current agent waypoints.

get_waypoints_mask()

Returns a functor of BxAxM boolean tensors representing current agent waypoints present mask.

get_all_agents_absolute()

Returns a functor of BxAx6 tensors, where the last dimension contains the following information: x, y, psi, length, width, present. Typically used to implement non-visual observation modalities.

get_all_agents_relative(exclude_self=True)

Returns a functor of BxAxAllx6 tensors, specifying for each of A agents the relative position about the other agents. ‘All’ is the number of all agents in the simulation, including hidden ones, across all agent types. If exclude_self is set, for each agent in A, that agent itself is removed from All. The final dimension has the same meaning as in get_all_agents_absolute, except now the positions and orientations are relative to the specified agent.

set_state(agent_state, mask=None)

Arbitrarily set the state of the agents, without advancing the simulation. The change is effective immediately, without waiting for the next step.

Parameters:

• agent_state – a functor of BxAx4 tensors with agent states

• mask – a functor of BxA boolean tensors, deciding which agent states to update; all by default

fit_action(future_state, current_state=None)

Computes an action that would (aproximately) produce the desired state.

Parameters:

• future_state – a functor of BxAx4 tensors defining the desired state

• current_state – if different from the current simulation state, in the same format as future state

Returns:

a functor of BxAxAc tensors

render(camera_xy, camera_psi, res=None, rendering_mask=None, fov=None, waypoints=None, waypoints_rendering_mask=None)

Renders the world from bird’s eye view using cameras in given positions.

Parameters:

• camera_xy – BxNx2 tensor of x-y positions for N cameras

• camera_psi – BxNx1 tensor of orientations for N cameras

• res – desired image resolution (only square resolutions are supported; by default use value from config)

• rendering_mask – functor of BxNxA tensors, indicating which agents should be rendered each camera

• fov – the field of view of the resulting image in meters (by default use value from config)

• waypoints – BxNxMx2 tensor of M waypoints per camera (x,y)

• waypoints_rendering_mask – BxNxM tensor of M waypoint masks per camera, indicating which waypoints should be rendered

Returns:

BxNxCxHxW tensor of resulting RGB images for each camera

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.

abstract update_present_mask(*args, **kwargs)

Sets the present mask of agents to the provided value.

Parameters:

present_mask – a functor of BxA boolean tensors

_compute_collision_of_multi_agents(mask=None)

Computes the collision metric for selected (default all) agents in the simulation. Includes collisions with all agents in the simulation, including the ones not exposed through the interface of this class. Used with nograd and nograd-pytorch3d metrics.

Parameters:

mask – a functor of BxA boolean tensors, indicating for which agents to compute the loss (by default use present mask)

Returns:

a functor of BxA tensors

class torchdrivesim.simulator.BoundedRegionWrapper(simulator: SimulatorInterface, exposed_agent_limit: int, default_action: TensorPerAgentType, warmup_timesteps: int, cutoff_polygon_verts: TensorPerAgentType)

Bases: SelectiveWrapper

A variant of SelectiveWrapper that exposes agents contained within a given polygon defining the area of interest. It additionally provides a method to determine which of the exposed agents have been “warmed-up”, which means that they have been exposed for a certain number of time steps, but that is separate from the main functionality.

Parameters:

• warmup_timesteps – after how many steps of being exposed are the agents considered “warmed-up”

• cutoff_polygon_verts – vertices defining the bounding convex polygon, in either clockwise or counter-closkwise order, provided as a functor of BxNx2 tensors, so that each agent type can use a different polygon

update_exposed_agents()

There should always be E exposed agents, although some of them may be marked as absent.

get_present_mask()

Returns a functor of BxA boolean tensors indicating which agents are currently present in the simulation.

is_warmed_up()

Returns:

a functor of BxA boolean tensors indicating which agents are warmed-up

to(device) → typing_extensions.Self

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

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.

extend(n, in_place=True)

Multiplies the first batch dimension by the given number. Like in pytorch3d, this is equivalent to introducing extra batch dimension on the right and then flattening.

select_batch_elements(idx, in_place=True)

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

class torchdrivesim.simulator.NoReentryBoundedRegionWrapper(simulator: SimulatorInterface, exposed_agent_limit, default_action, warmup_timesteps, cutoff_polygon_verts: TensorPerAgentType | None = None)

Bases: BoundedRegionWrapper

A variant of BoundedRegionWrapper that does not allow reentry. This means that if an agent is not exposed at a certain time, it will never be exposed in the future. However, the agents that can stop being exposed if they exit the area of interest.

to(device) → typing_extensions.Self

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

extend(n, in_place=True)

Multiplies the first batch dimension by the given number. Like in pytorch3d, this is equivalent to introducing extra batch dimension on the right and then flattening.

select_batch_elements(idx, in_place=True)

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

update_exposed_agents()

There should always be E exposed agents, although some of them may be marked as absent.

update_present_mask(present_mask)

Sets the present mask of agents to the provided value.

Parameters:

present_mask – a functor of BxA boolean tensors

get_present_mask()

Returns a functor of BxA boolean tensors indicating which agents are currently present in the simulation.