from dataclasses import dataclass from typing import Callable, Dict, List, Optional, Tuple from mmcv import Config import numpy as np import torch from risk_biased.mpc_planner.dynamics import PositionVelocityDoubleIntegrator from risk_biased.mpc_planner.planner_cost import TrackingCost from risk_biased.predictors.biased_predictor import LitTrajectoryPredictor from risk_biased.utils.cost import BaseCostTorch from risk_biased.utils.planner_utils import ( AbstractState, to_state, evaluate_risk, get_interaction_cost, ) from risk_biased.utils.risk import AbstractMonteCarloRiskEstimator @dataclass class CrossEntropySolverParams: """Dataclass for Cross Entropy Solver Parameters Args: num_control_samples: number of Monte Carlo samples for control input num_elite: number of elite samples iter_max: maximum iteration number smoothing_factor: smoothing factor in (0, 1) used to update the mean and the std of the control input distribution for the next iteration. If 0, the updated distribution is independent of the previous iteration. If 1, the updated distribution is the same as the previous iteration. mean_warm_start: internally saves control_input_mean of the last iteration of the current solve, so that control_input_mean will be warm-started in the next solve """ num_control_samples: int num_elite: int iter_max: int smoothing_factor: float mean_warm_start: bool dt: float @staticmethod def from_config(cfg: Config): return CrossEntropySolverParams( cfg.num_control_samples, cfg.num_elite, cfg.iter_max, cfg.smoothing_factor, cfg.mean_warm_start, cfg.dt, ) class CrossEntropySolver: """Cross Entropy Solver for MPC Planner Args: params: CrossEntropySolverParams object dynamics_model: dynamics model for control control_input_mean: (num_agents, num_steps_future, control_dim) tensor of control input mean control_input_std: (num_agents, num_steps_future, control_dim) tensor of control input std tracking_cost_function: deterministic tracking cost that does not involve ado intraction_cost_function: interaction cost function between ego and (stochastic) ado risk_estimator (optional): Monte Carlo risk estimator for risk computation. If None, risk-neutral expecation is used for selectoin of elites. Defaults to None. """ def __init__( self, params: CrossEntropySolverParams, dynamics_model: PositionVelocityDoubleIntegrator, control_input_mean: torch.Tensor, control_input_std: torch.Tensor, tracking_cost_function: TrackingCost, interaction_cost_function: BaseCostTorch, risk_estimator: Optional[AbstractMonteCarloRiskEstimator] = None, ) -> None: self.params = params self.control_input_mean_init = control_input_mean.detach().clone() self.control_input_std_init = control_input_std.detach().clone() assert ( self.control_input_mean_init.shape == self.control_input_std_init.shape ), "control input mean and std must have the same size" assert ( self.control_input_mean_init.shape[-1] == dynamics_model.control_dim ), f"control dimension must be {dynamics_model.control_dim}" self.dynamics_model = dynamics_model self.tracking_cost = tracking_cost_function self.interaction_cost = interaction_cost_function self.risk_estimator = risk_estimator self._iter_current = None self._control_input_mean = None self._control_input_std = None self._latest_ado_position_future_samples = None self.reset() def reset(self) -> None: """Resets the solver's internal state""" self._iter_current = 0 self._control_input_mean = self.control_input_mean_init.clone() self._control_input_std = self.control_input_std_init.clone() self._latest_ado_position_future_samples = None def step( self, ego_state_history: AbstractState, ego_state_target_trajectory: AbstractState, ado_state_future_samples: AbstractState, weights: torch.Tensor, verbose: bool = False, risk_level: float = 0.0, ) -> Dict: """Performs one iteration step of the Cross Entropy Method Args: ego_state_history: (num_agents, num_steps) ego state history ego_state_target_trajectory: (num_agents, num_steps_future) ego target state trajectory ado_state_future_samples: (num_prediction_samples, num_agents, num_steps_future) predicted ado trajectory samples weights: (num_prediction_samples, num_agents) prediction sample weight verbose (optional): Print progress. Defaults to False. risk_level (optional): a risk-level float for the solver. If 0.0, risk-neutral expectation is used for selection of elites. Defaults to 0.0. Return: Dictionary containing information about this solver step. """ self._iter_current += 1 ego_control_input = torch.normal( self._control_input_mean.expand( self.params.num_control_samples, -1, -1, -1 ), self._control_input_std.expand(self.params.num_control_samples, -1, -1, -1), ) if verbose: print(f"**Cross Entropy Iteration {self._iter_current}") print( f"****Drawring ego's control input samples of {ego_control_input.shape}" ) ego_state_current = ego_state_history[..., -1] ego_state_future = self.dynamics_model.simulate( ego_state_current, ego_control_input ) if verbose: print(f"****Simulating ego's future state trajectory") # state starts with x, y, angle, vx, vy tracking_cost = self.tracking_cost( ego_state_future.position, ego_state_target_trajectory.position, ego_state_target_trajectory.velocity, ) if verbose: print( f"****Computing tracking cost of {tracking_cost.shape} for the control input samples" ) # state starts with x, y interaction_cost = get_interaction_cost( ego_state_future, ado_state_future_samples, self.interaction_cost, ) if verbose: print( f"****Computing interaction cost of {interaction_cost.shape} for the control input samples" ) interaction_risk = evaluate_risk( risk_level, interaction_cost, weights.permute(1, 0).unsqueeze(0).expand_as(interaction_cost), self.risk_estimator, ) total_risk = interaction_risk + tracking_cost elite_ego_control_input, elite_total_risk = self._get_elites( ego_control_input, total_risk ) if verbose: print(f"****Selecting {self.params.num_elite} elite samples") print(f"****Elite Total_Risk Information: {elite_total_risk}") info = dict( iteration=self._iter_current, control_input_mean=self._control_input_mean.detach().cpu().numpy().copy(), control_input_std=self._control_input_std.detach().cpu().numpy().copy(), ego_state_future=ego_state_future.get_states(5) .detach() .cpu() .numpy() .copy(), ado_state_future_samples=ado_state_future_samples.get_states(5) .detach() .cpu() .numpy() .copy(), sample_weights=weights.detach().cpu().numpy().copy(), tracking_cost=tracking_cost.detach().cpu().numpy().copy(), interaction_cost=interaction_cost.detach().cpu().numpy().copy(), total_risk=total_risk.detach().cpu().numpy().copy(), ) self._update_control_distribution(elite_ego_control_input) if verbose: print("****Updating ego's control distribution") return info def solve( self, predictor: LitTrajectoryPredictor, ego_state_history: AbstractState, ego_state_target_trajectory: AbstractState, ado_state_history: AbstractState, normalizer: Callable[[torch.Tensor], Tuple[torch.Tensor, torch.Tensor]], num_prediction_samples: int = 1, verbose: bool = False, risk_level: float = 0.0, resample_prediction: bool = False, risk_in_predictor: bool = False, ) -> List[Dict]: """Performs Cross Entropy optimization of ego's control input Args: predictor: LitTrajectoryPredictor object ego_state_history: (num_agents, num_steps, state_dim) ego state history ego_state_target_trajectory: (num_agents, num_steps_future, state_dim) ego target state trajectory ado_state_history: (num_agents, num_steps, state_dim) ado state history normalizer: function that takes in an unnormalized trajectory and that outputs the normalized trajectory and the offset in this order num_prediction_samples: number of prediction samples. Defaults to 1. verbose (optional): Print progress. Defaults to False. risk_level (optional): a risk-level float for the entire prediction-planning pipeline. If 0.0, risk-neutral prediction and planning are used. Defaults to 0.0. resample_prediction (optional): If True, prediction is re-sampled in each cross-entropy iteration. Defaults to False. risk_in_predictor (optional): If True, risk-biased prediction is used and the solver becomes risk-neutral. If False, risk-neutral prediction is used and the solver becomes risk-sensitive. Defaults to False. Return: List of dictionaries each containing information about the corresponding solver step. """ if risk_level == 0.0: risk_level_planner, risk_level_predictor = 0.0, 0.0 else: if risk_in_predictor: risk_level_planner, risk_level_predictor = 0.0, risk_level else: risk_level_planner, risk_level_predictor = risk_level, 0.0 self.reset() infos = [] ego_state_future = self.dynamics_model.simulate( ego_state_history[..., -1], self.control_sequence, ) for iter in range(self.params.iter_max): assert iter == self._iter_current if resample_prediction or self._iter_current == 0: ado_state_future_samples, weights = self.sample_prediction( predictor, ado_state_history, normalizer, ego_state_history, ego_state_future, num_prediction_samples, risk_level_predictor, ) self._latest_ado_position_future_samples = ado_state_future_samples info = self.step( ego_state_history, ego_state_target_trajectory, ado_state_future_samples, weights, verbose=verbose, risk_level=risk_level_planner, ) infos.append(info) if self.params.mean_warm_start: self.control_input_mean_init[:, :-1] = ( self._control_input_mean[:, 1:].detach().clone() ) return infos @property def control_sequence(self) -> torch.Tensor: """Returns the planned control sequence, which is a detached copy of the control input mean tensor Returns: (num_steps_future, control_dim) control sequence tensor """ return self._control_input_mean.detach().clone() @staticmethod def sample_prediction( predictor: LitTrajectoryPredictor, ado_state_history: AbstractState, normalizer: Callable[[torch.Tensor], Tuple[torch.Tensor, torch.Tensor]], ego_state_history: AbstractState, ego_state_future: AbstractState, num_prediction_samples: int = 1, risk_level: float = 0.0, ) -> Tuple[AbstractState, torch.Tensor]: """Sample prediction from the predictor given the history, normalizer, and the desired risk-level Args: predictor: LitTrajectoryPredictor object ado_state_history: (num_agents, num_steps, state_dim) tensor of ado position history normalizer: function that takes in an unnormalized trajectory and that outputs the normalized trajectory and the offset in this order ego_state_history: (num_agents, num_steps , state_dim) tensor of ego position history or future ego_state_future: (num_agents, num_steps_future, state_dim) tensor of ego position history or future num_prediction_samples (optional): number of prediction samples. Defaults to 1. risk_level (optional): a risk-level float for the predictor. If 0.0, risk-neutral prediction is sampled. Defaults to 0.0. Returns: state samples of shape (num_agents, num_prediction_samples, num_steps_future) probability weights of the samples of shape (num_agents, num_prediction_samples) """ ado_position_history_normalized, offset = normalizer( ado_state_history.get_states(predictor.dynamic_state_dim) .unsqueeze(0) .expand(num_prediction_samples, -1, -1, -1) ) x = ado_position_history_normalized.clone() mask_x = torch.ones_like(x[..., 0]) map = torch.empty(num_prediction_samples, 0, 0, 2, device=x.device) mask_map = torch.empty(num_prediction_samples, 0, 0, device=x.device) batch = ( x, mask_x, map, mask_map, offset, ego_state_history.get_states(predictor.dynamic_state_dim) .unsqueeze(0) .expand(num_prediction_samples, -1, -1, -1), ego_state_future.get_states(predictor.dynamic_state_dim) .unsqueeze(0) .expand(num_prediction_samples, -1, -1, -1), ) ado_position_future_samples, weights = predictor.predict_step( batch, 0, risk_level=risk_level, return_weights=True, ) ado_position_future_samples = ado_position_future_samples.detach().cpu() weights = weights.detach().cpu() return to_state(ado_position_future_samples, predictor.dt), weights def fetch_latest_prediction(self): if self._latest_ado_position_future_samples is not None: return self._latest_ado_position_future_samples else: return None def _get_elites( self, control_input: torch.Tensor, risk: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: """Selects elite control input based on corresponding risk (lower the better) Args: control_input: (num_control_samples, num_agents, num_steps_future, control_dim) control samples risk: (num_control_samples, num_agents) risk tensor Returns: elite_control_input: (num_elite, num_agents, num_steps_future, control_dim) elite control elite_risk: (num_elite, num_agents) elite risk """ num_control_samples = self.params.num_control_samples assert ( control_input.shape[0] == num_control_samples ), f"size of control_input tensor must be {num_control_samples} at dimension 0" assert ( risk.shape[0] == num_control_samples ), f"size of risk tensor must be {num_control_samples} at dimension 0" _, sorted_risk_indices = torch.sort(risk, dim=0) elite_control_input = control_input[ sorted_risk_indices[: self.params.num_elite], np.arange(risk.shape[1]) ] elite_risk = risk[ sorted_risk_indices[: self.params.num_elite], np.arange(risk.shape[1]) ] return elite_control_input, elite_risk def _update_control_distribution(self, elite_control_input: torch.Tensor) -> None: """Updates control input distribution using elites Args: elite_control_input: (num_elite, num_steps_future, control_dim) elite control """ num_elite, smoothing_factor = ( self.params.num_elite, self.params.smoothing_factor, ) assert ( elite_control_input.shape[0] == num_elite ), f"size of elite_control_input tensor must be {num_elite} at dimension 0" elite_control_input_mean = elite_control_input.mean(dim=0, keepdim=False) if num_elite < 2: elite_control_input_std = torch.zeros_like(elite_control_input_mean) else: elite_control_input_std = elite_control_input.std(dim=0, keepdim=False) self._control_input_mean = ( 1.0 - smoothing_factor ) * elite_control_input_mean + smoothing_factor * self._control_input_mean self._control_input_std = ( 1.0 - smoothing_factor ) * elite_control_input_std + smoothing_factor * self._control_input_std