Spaces:
Running
Running
File size: 1,608 Bytes
5769ee4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
import torch
from risk_biased.utils.planner_utils import AbstractState, to_state
class PositionVelocityDoubleIntegrator:
"""Deterministic discrete-time double-integrator dynamics, where state is
[position_x_m, position_y_m, velocity_x_m_s velocity_y_m_s] and control is
[acceleration_x_m_s2, acceleration_y_m_s2].
Args:
dt: time differential between two discrete timesteps in seconds
"""
def __init__(self, dt: float):
self.dt = dt
self.control_dim = 2
def simulate(
self,
state_init: AbstractState,
control_input: torch.Tensor,
) -> AbstractState:
"""Euler-integrate dynamics from the initial position and the initial velocity given
an acceleration input
Args:
state_init: initial Markov state of the system
control_input: (num_agents, num_steps_future, 2) tensor of acceleration input
Returns:
(num_agents, num_steps_future, 5) tensor of simulated future Markov state
sequence
"""
position_init, velocity_init = state_init.position, state_init.velocity
assert (
control_input.shape[-1] == self.control_dim
), "invalid control input dimension"
velocity_future = velocity_init + self.dt * torch.cumsum(control_input, dim=-2)
position_future = position_init + self.dt * torch.cumsum(
velocity_future, dim=-2
)
state_future = to_state(
torch.cat((position_future, velocity_future), dim=-1), self.dt
)
return state_future
|