Code for the paper "Deep Model Predictive Optimization" 
import torch
import torch.nn as nn
import torch.distributions as D

from ..models.net_utils import create_net
from ..mpc.task.base_rollout_task import BaseRolloutTask
from .. import utils

from torch import Tensor
from torch.nn import Module
from typing import Dict, Any, Optional, List, Union, Tuple, Callable

STD_MAX = 1e3
STD_MIN = 1e-6

class DMPOPolicy(nn.Module):
    """
    Contains actor and critic for DMPO
    """
    def __init__(self,
                 d_action: int,
                 d_state: int,
                 actor_params: Dict[str, Any],
                 critic_params: Dict[str, Any],
                 shift_params: Optional[Dict[str, Any]] = None,
                 rollout_task: Optional[BaseRolloutTask] = None,
                 horizon: int = 1,
                 num_particles: int = 1,
                 gamma: float = 1.,
                 top_k: int = 1,
                 n_iters: int = 1,
                 init_mean: Union[float, Tensor] = 0.,
                 init_std: Union[float, Tensor] = 1.,
                 sample_params: Optional[Dict[str, Any]] = None,
                 seed_val: int = 0,
                 state_seq_key: str = 'state_seq',
                 mppi_params: Optional[Dict[str, Any]] = None,
                 d_cond: Optional[int] = None,
                 cond_mode: Optional[str] = None,
                 cond_actor: bool = False,
                 cond_critic: bool = False,
                 cond_shift: bool = False,
                 critic_use_cost: bool = False,
                 actor_use_state: bool = False,
                 use_mean: bool = False,
                 mppi_mode: bool = False,
                 is_delta: bool = False,
                 is_gated: bool = False,
                 is_residual: bool = False,
                 mean_search_std: Optional[float] = None,
                 std_search_std: Optional[float] = None,
                 learn_search_std: bool = False,
                 learn_rollout_std: bool = False,
                 action_lows: Optional[Union[float, List[float], Tensor]] = None,
                 action_highs: Optional[Union[float, List[float], Tensor]] = None,
                 state_scale: Optional[Union[float, List[float], Tensor]] = None,
                 cond_scale: Optional[Union[float, List[float], Tensor]] = None,
                 tensor_args: Dict[str, Any] = {'device': 'cpu', 'dtype': torch.float32}):
        """
        :param d_action: Dimensionality of the action space
        :param d_state: Dimensionality of the state space
        :param actor_params: Dictionary of model parameters for the actor network
        :param critic_params: Dictionary of model parameters for the critic network
        :param shift_params: Dictionary of model parameters for the shift model network
        :param rollout_task: Rollout function
        :param horizon: Horizon of MPC controller
        :param num_particles: Number of particles used for rollouts
        :param gamma: Discount factor
        :param top_k: Number of top performing trajectories to save for visualization purposes
        :param n_iters: Number of optimizer iterations
        :param init_mean: Initial mean of the optimizee
        :param init_std: Initial STD of the optimizee
        :param sample_params: Dictionary of parameters for sampling
        :param seed_val: Seed used to generate fixed samples from the standard Gaussian
        :param state_seq_key: Key used to acquire state sequences from dictionary returned by rollout module
        :param mppi_params: Dictionary of MPPI parameters
        :param d_cond: Dimensionality of conditioning variable
        :param cond_mode: Method for incorporating the conditioning variable (e.g. concatenating it to network input)
        :param cond_actor: Flag for conditioning the actor
        :param cond_critic: Flag for conditioning the critic
        :param cond_shift: Flag for conditioning the shift model
        :param critic_use_cost: Flag to indicate that the critic should use the trajectory costs rather than state
        :param actor_use_state: Flag to indicate that the actor should also use state
        :param use_mean: Flag to indicate that we use the mean of our optimizer distribution (for test mode)
        :param mppi_mode: Flag to indicate that we should run vanilla MPPI
        :param is_delta: Flag to indicate that the actor output should be added to current plan (used if not a residual on MPPI)
        :param is_gated: Flag to indicate that we use a gating term on the mean update
        :param is_residual: Flag to indicate that DMPO is learning residuals on MPPI
        :param mean_search_std: Initial STD of the optimizer policy over optimizee means
        :param std_search_std: Initial STD of the optimizer policy over optimizee STDs
        :param learn_search_std: Flag to indicate that we should learn the search STDs
        :param learn_rollout_std: Flag to indicate that we should learn the optimizee STDs (rather than assuming they are fixed)
        :param action_lows: Minimum action values
        :param action_highs: Maximum action values
        :param state_scale: Scale factor on state prior to any network input
        :param cond_scale: Scale factor on conditioning variable prior to any network input
        :param tensor_args: PyTorch Tensor settings
        """
        super().__init__()

        # Rollout parameters
        self.d_action = d_action
        self.d_state = d_state
        self.d_cond = d_cond
        self.horizon = horizon

        self.num_particles = num_particles
        self.gamma = gamma
        self.top_k = top_k
        self.n_iters = n_iters
        self.init_n_iters = n_iters
        self.seed_val = seed_val
        self.state_seq_key = state_seq_key

        # Rollout function
        self.rollout_task = rollout_task

        # Sample parameters
        self.use_halton = sample_params.get('use_halton', True) if not sample_params is None else True

        # MPPI parameters
        self.temperature = mppi_params.get('temperature', 1e-3) if not mppi_params is None else 1e-3
        self.step_size = mppi_params.get('step_size', 1) if not mppi_params is None else 1
        self.scale_costs = mppi_params.get('scale_costs', True) if not mppi_params is None else True

        # Conditioning parameters
        self.cond_mode = cond_mode if not self.d_cond is None else None
        self.cond_actor = cond_actor
        self.cond_critic = cond_critic
        self.cond_shift = cond_shift

        # Other parameters
        self.use_mean = use_mean
        self.mppi_mode = mppi_mode
        self.learn_search_std = learn_search_std
        self.learn_rollout_std = learn_rollout_std
        self.is_gated = is_gated
        self.is_residual = is_residual
        self.is_delta = is_delta
        self.critic_use_cost = critic_use_cost
        self.actor_use_state = actor_use_state

        self.action_lows = utils.to_tensor(action_lows, tensor_args) if not action_lows is None else None
        self.action_highs = utils.to_tensor(action_highs, tensor_args) if not action_highs is None else None
        self.state_scale = utils.to_tensor(state_scale, tensor_args) if not state_scale is None else None
        self.cond_scale = utils.to_tensor(cond_scale, tensor_args) if not cond_scale is None else None
        self.tensor_args = tensor_args

        # Sizes of each input type
        mean_size = self.horizon * self.d_action
        std_size = self.horizon * self.d_action
        cost_size = self.num_particles
        cond_size = self.d_cond

        # Create the shift network
        if not shift_params is None:
            in_size = mean_size
            if learn_rollout_std:
                in_size += std_size
            if cond_mode == 'cat' and not d_cond is None and cond_shift:
                in_size += cond_size

            out_size = self.horizon*self.d_action
            if learn_rollout_std:
                out_size += self.horizon*self.d_action

            self.shift_model = create_net(in_size=in_size, out_size=out_size, **shift_params)
        else:
            self.shift_model = None

        # Create the actor network
        in_size = mean_size + cost_size
        if actor_use_state:
            in_size += d_state
        if learn_rollout_std:
            in_size += std_size
        if cond_mode == 'cat' and not d_cond is None and cond_actor:
            in_size += cond_size

        out_size = mean_size
        if is_gated:
            out_size += mean_size
        if learn_search_std:
            out_size += mean_size
        if learn_rollout_std:
            out_size += std_size
            if learn_search_std:
                out_size += std_size

        self.actor = create_net(in_size=in_size, out_size=out_size, **actor_params)

        # Create the critic
        if not critic_use_cost:
            in_size = d_state + mean_size
        else:
            in_size = mean_size + cost_size

        if learn_rollout_std:
            in_size += std_size
        if cond_mode == 'cat' and not d_cond is None and cond_critic:
            in_size += cond_size

        self.critic = create_net(in_size=in_size, out_size=1, **critic_params)

        # Set the initial mean
        if not isinstance(init_mean, Tensor):
            if isinstance(init_mean, float):
                self.init_mean = init_mean*torch.ones((horizon, d_action), **tensor_args)
            else:
                self.init_mean = torch.tensor(init_mean, **tensor_args).unsqueeze(0).repeat(horizon, 1)
        else:
            self.init_mean = init_mean.to(**tensor_args)

        # Set the initial rollout STD
        if not isinstance(init_std, Tensor):
            if isinstance(init_std, float):
                self.init_std = init_std*torch.ones((horizon, d_action), **tensor_args)
            else:
                self.init_std = torch.tensor(init_std, **tensor_args).unsqueeze(0).repeat(horizon, 1)
        else:
            self.init_std = init_std.to(**tensor_args)

        # Set the initial mean search STD
        if not mean_search_std is None:
            if not isinstance(mean_search_std, Tensor):
                if isinstance(mean_search_std, float):
                    self.mean_search_std = mean_search_std * torch.ones((horizon, d_action), **tensor_args)
                else:
                    self.mean_search_std = torch.tensor(mean_search_std, **tensor_args).unsqueeze(0).repeat(horizon,
                                                                                                            1)
            else:
                self.mean_search_std = mean_search_std.to(**tensor_args)
        else:
            self.mean_search_std = None

        # Set the initial STD search STD
        if not std_search_std is None:
            if not isinstance(std_search_std, Tensor):
                if isinstance(std_search_std, float):
                    self.std_search_std = std_search_std * torch.ones((horizon, d_action), **tensor_args)
                else:
                    self.std_search_std = torch.tensor(std_search_std, **tensor_args).unsqueeze(0).repeat(horizon,
                                                                                                          1)
            else:
                self.std_search_std = std_search_std.to(**tensor_args)
        else:
            self.std_search_std = None

    def set_n_iters(self, n_iters: int):
        self.n_iters = n_iters

    def reset(self):
        self.mean = self.init_mean.clone()
        self.std = self.init_std.clone()
        self.samples = None
        self.state_samples = None
        self.mppi_mean = None
        self.params_stacked = {}

    def update_params(self, kwargs: Dict[str, Any]):
        kwargs_stacked = {key: [] for key in kwargs[0].keys()}
        for kwargs_dict in kwargs:
            for k, v in kwargs_dict.items():
                kwargs_stacked[k].append(v)
        self.params_stacked = kwargs_stacked
        self.rollout_task.update_params(kwargs_stacked)

    def set_task(self, task: BaseRolloutTask) -> None:
        self.rollout_task = task

    def generate_samples(self, batch_size: int):
        num_particles = self.num_particles-1

        if self.samples is None:
            if self.use_halton:
                self.samples = utils.generate_gaussian_halton_samples(num_samples=num_particles,
                                                                      ndims=self.d_action*self.horizon,
                                                                      seed_val=self.seed_val,
                                                                      device=self.tensor_args['device'],
                                                                      dtype=self.tensor_args['dtype'])
            else:
                with torch.random.fork_rng([torch.device(self.tensor_args['device'])]) as rng:
                    torch.random.manual_seed(self.seed_val)
                    self.samples = torch.randn((num_particles, self.d_action*self.horizon),
                                               **self.tensor_args)

            self.samples = self.samples[None, :, :]

        samples = self.samples.view(1, num_particles, self.horizon, self.d_action)
        samples = samples.repeat(batch_size, 1, 1, 1)

        # Always ensure mean is a sample
        zeros = torch.zeros((batch_size, 1, self.horizon, self.d_action), **self.tensor_args)
        samples = torch.cat((zeros, samples), dim=1)

        std = self.std[:, None, :, :]
        samples = samples * std + self.mean[:, None, :, :]

        if not self.action_lows is None and not self.action_highs is None:
            samples = samples.clamp(self.action_lows, self.action_highs)
        return samples

    def run_rollouts(self, x: Tensor) -> Tuple[Tensor, Tensor]:
        batch_size = x.shape[0]

        with torch.no_grad():
            # Generate action samples
            samples = self.generate_samples(batch_size)

            # Run the rollouts
            trajectories = self.rollout_task.run_rollouts(x, samples)

            # Collect the results
            costs = trajectories['costs']
            states = trajectories[self.state_seq_key]

            if not isinstance(costs, Tensor):
                costs = costs['cost']

            # Compute the total costs
            gamma = torch.tensor([self.gamma ** i for i in range(costs.shape[-1])], **self.tensor_args)
            costs = torch.sum(gamma[None, None, :] * costs, dim=-1)

            # Process states
            self.state_samples = states
            top_values, top_idx = torch.topk(-costs, self.top_k, dim=-1)
            self.top_values = -top_values
            self.top_idx = top_idx
            self.top_trajs = [torch.index_select(states[idx], 0, top_idx[idx]) for idx in range(states.shape[0])]

        return costs, samples

    def process_mean(self, x: Tensor, mean: Tensor):
        batch_size = x.shape[0]
        if mean.ndim == 2:
            mean = mean.unsqueeze(0).repeat(batch_size, 1, 1)
        old_mean = mean

        # Shift mean forward
        shifted_mean = torch.cat((mean[:, 1:], torch.zeros_like(mean[:, -1:])), dim=1)
        return shifted_mean, old_mean

    def process_std(self, x: Tensor, std: Tensor, shift=True):
        batch_size = x.shape[0]
        if std.ndim == 2:
            std = std.unsqueeze(0).repeat(batch_size, 1, 1)
        old_std = std

        # Shift mean forward
        if shift:
            shifted_std = torch.cat((std[:, 1:], std[:, -1:]), dim=1)
            return shifted_std, old_std
        else:
            return old_std

    def get_actor_embedding(self,
                            x: Tensor,
                            costs: Tensor,
                            mean: Optional[Tensor] = None,
                            std: Optional[Tensor] = None,
                            cond: Optional[Tensor] = None) -> Tensor:
        batch_size = x.shape[0]
        if not self.state_scale is None:
           x = x/self.state_scale

        # Compute the normalized costs
        costs = costs.reshape(batch_size, -1)
        cost_mean = costs.mean(dim=-1)
        cost_std = costs.std(dim=-1)
        costs = (costs - cost_mean[:, None]) / (cost_std[:, None] + 1e-6)

        # Reshape the mean and STD
        if not self.action_highs is None and not self.action_lows is None:
           mean = (mean - self.action_lows) / (self.action_highs - self.action_lows)
        mean = mean.view(batch_size, -1)

        if not std is None:
            if not self.action_highs is None and not self.action_lows is None:
                std = std / (self.init_std[0]*10)
            std = std.view(batch_size, -1)

        # Prepare the condition
        if not cond is None and not self.cond_scale is None:
            if self.cond_scale != cond.shape[-1]:
                cond_scale = self.cond_scale[:1].repeat(cond.shape[-1])
                cond = cond / cond_scale
            else:
                cond = cond / self.cond_scale

        # Form the network input
        net_in = torch.cat((costs, mean), dim=-1)
        if self.actor_use_state:
            net_in = torch.cat((net_in, x), dim=-1)
        if self.learn_rollout_std:
            net_in = torch.cat((net_in, std), dim=-1)
        if not cond is None and not self.cond_mode is None and self.cond_actor:
            if self.cond_mode == 'cat':
                net_in = torch.cat((net_in, cond), dim=-1)
            else:
                raise ValueError('Invalid condition mode {} specified.'.format(self.cond_mode))

        return net_in

    def get_critic_embedding(self,
                            x: Tensor,
                            costs: Tensor,
                            mean: Optional[Tensor] = None,
                            std: Optional[Tensor] = None,
                            cond: Optional[Tensor] = None) -> Tensor:

        batch_size = x.shape[0]
        if not self.state_scale is None:
           x = x/self.state_scale

        # Reshape the mean and STD
        if not self.action_highs is None and not self.action_lows is None:
           mean = (mean - self.action_lows) / (self.action_highs - self.action_lows)
        mean = mean.view(batch_size, -1)

        if not std is None:
            if not self.action_highs is None and not self.action_lows is None:
                std = std / (self.init_std[0]*10)
            std = std.view(batch_size, -1)

        # Compute the total costs
        costs = costs.reshape(batch_size, -1)
        cost_mean = costs.mean(dim=-1)
        cost_std = costs.std(dim=-1)
        costs = (costs - cost_mean[:, None]) / (cost_std[:, None] + 1e-6)

        # Prepare the condition
        if not cond is None and not self.cond_scale is None:
            if self.cond_scale != cond.shape[-1]:
                cond_scale = self.cond_scale[:1].repeat(cond.shape[-1])
                cond = cond / cond_scale
            else:
                cond = cond / self.cond_scale

        # Form the network input
        if not self.critic_use_cost:
            net_in = torch.cat((x, mean), dim=-1)
        else:
            net_in = torch.cat((costs, mean), dim=-1)

        if self.learn_rollout_std:
            net_in = torch.cat((net_in, std), dim=-1)

        if not cond is None and not self.cond_mode is None and self.cond_critic:
            if self.cond_mode == 'cat':
                net_in = torch.cat((net_in, cond), dim=-1)
            else:
                raise ValueError('Invalid condition mode {} specified.'.format(self.cond_mode))

        return net_in

    def get_shift_embedding(self,
                            x: Tensor,
                            mean: Optional[Tensor] = None,
                            std: Optional[Tensor] = None,
                            shifted_std: Optional[Tensor] = None,
                            cond: Optional[Tensor] = None) -> Tensor:
        batch_size = x.shape[0]

        # Reshape the mean and STD
        if not self.action_highs is None and not self.action_lows is None:
           mean = (mean - self.action_lows) / (self.action_highs - self.action_lows)
        mean = mean.view(batch_size, -1)

        if not std is None:
            if not shifted_std is None:
                std = std / (shifted_std*10)
            std = std.view(batch_size, -1)

        # Prepare the condition
        if not cond is None and not self.cond_scale is None:
            if self.cond_scale != cond.shape[-1]:
                cond_scale = self.cond_scale[:1].repeat(cond.shape[-1])
                cond = cond / cond_scale
            else:
                cond = cond / self.cond_scale

        # Form the network input
        if self.learn_rollout_std:
            net_in = torch.cat((mean, std), dim=-1)
        else:
            net_in = mean

        if not cond is None and not self.cond_mode is None and self.cond_actor:
            if self.cond_mode == 'cat':
                net_in = torch.cat((net_in, cond), dim=-1)
            else:
                raise ValueError('Invalid condition mode {} specified.'.format(self.cond_mode))

        return net_in

    def run_actor(self,
                  x: Tensor,
                  costs: Tensor,
                  mean: Optional[Tensor] = None,
                  std: Optional[Tensor] = None,
                  cond: Optional[Tensor] = None) -> Tuple[D.Distribution, D.Distribution]:

        # Compute the MPPI update
        if self.mppi_mode or self.is_residual:
            with torch.no_grad():
                mppi_costs = costs
                if self.scale_costs:
                    min_costs = mppi_costs.min(dim=-1)[0][:, None]
                    max_costs = mppi_costs.max(dim=-1)[0][:, None]
                    mppi_costs = (mppi_costs - min_costs)/(max_costs - min_costs + 1e-6)

                weights = torch.softmax(-mppi_costs / self.temperature, dim=1)
                samples = self.generate_samples(weights.shape[0])
                update = torch.sum(samples * weights[:, :, None, None], dim=1)

            mppi_mean = (1 - self.step_size) * self.mean + self.step_size * update.view(-1, self.horizon, self.d_action)
        else:
            mppi_mean = None
        self.mppi_mean = mppi_mean

        # Compute the DMPO update
        if not self.mppi_mode:

            # Get the actor embedding
            net_in = self.get_actor_embedding(x, costs=costs, mean=mean, std=std, cond=cond)

            # Run the actor
            out = self.actor(net_in)
            mean = out[:, :self.horizon*self.d_action].view(-1, self.horizon, self.d_action)
            idx = 1

            # Scale the actor output so that it is in a reasonable range
            if not self.action_lows is None and not self.action_highs is None:
                mean = torch.tanh(mean)
                mean = mean*(self.action_highs - self.action_lows)

            # Get the gating term if used
            if self.is_gated:
                gating = out[:, self.horizon*self.d_action*idx:self.horizon*self.d_action*(idx+1)]
                gating = torch.tanh(gating).view(-1, self.horizon, self.d_action)
                idx += 1

            # Get the updated search STD for the optimizee mean
            if self.learn_search_std:
                log_mean_std = out[:, self.horizon*self.d_action*idx:self.horizon*self.d_action*(idx+1)]
                mean_std = log_mean_std.exp().view(-1, self.horizon, self.d_action)
                mean_std = torch.clamp(self.mean_search_std[None, :, :]*mean_std, STD_MIN, STD_MAX)
                idx += 1
            else:
                mean_std = self.mean_search_std

            # Handle updates to the optimizee STD if used
            if self.learn_rollout_std:

                # Compute the updated optimizee STD
                log_std = out[:, self.horizon*self.d_action*idx:self.horizon*self.d_action*(idx+1)]
                std = log_std.exp().view(-1, self.horizon, self.d_action)
                std = torch.clamp(self.init_std[None, :, :]*std, STD_MIN, STD_MAX)
                idx += 1

                # Get the updated search STD for the optimizee STD
                if self.learn_search_std:
                    log_std_std = out[:, self.horizon*self.d_action*idx:self.horizon*self.d_action*(idx+1)]
                    std_std = log_std_std.exp().view(-1, self.horizon, self.d_action)
                    std_std = torch.clamp(self.std_search_std*std_std, STD_MIN, STD_MAX)
                else:
                    std_std = self.std_search_std

            # Compute the updated mean
            if not self.is_residual and self.is_delta:
                # If not in residual mode and learning a delta on current plan, optionally with a gating term
                if self.is_gated:
                    mean = (1-gating)*self.mean + gating*mean
                else:
                    mean = mean + self.mean
            elif self.is_residual:
                # If in residual mode, form the update on the MPPI proposed mean, optionally using a gating term
                if self.is_gated:
                    mean = (1-gating)*mppi_mean + gating*mean
                else:
                    mean = mppi_mean + mean

            mean_dist = D.Normal(mean, mean_std)
            if self.learn_rollout_std:
                std_dist = D.Normal(std, std_std)
            else:
                std_dist = None
        else:
            mean_dist = D.Normal(mppi_mean, self.std)
            std_dist = None

        return mean_dist, std_dist

    def run_critic(self,
                   x: Tensor,
                   costs: Tensor,
                   mean: Optional[Tensor] = None,
                   std: Optional[Tensor] = None,
                   cond: Optional[Tensor] = None) -> Tensor:

        # Get the critic embedding
        net_in = self.get_critic_embedding(x, costs=costs, mean=mean, std=std, cond=cond)

        # Run the critics
        return self.critic(net_in)

    def run_shift_model(self,
                        x: Tensor,
                        mean: Optional[Tensor] = None,
                        shifted_mean: Optional[Tensor] = None,
                        std: Optional[Tensor] = None,
                        shifted_std: Optional[Tensor] = None,
                        cond: Optional[Tensor] = None) -> Tuple[Tensor, Tensor]:
        # Get the shift model embedding
        net_in = self.get_shift_embedding(x, mean=mean, std=std, shifted_std=shifted_std, cond=cond)

        # Run the shift model
        out = self.shift_model(net_in)

        if self.learn_rollout_std:
            mean, log_std = torch.split(out, self.horizon*self.d_action, dim=-1)
            mean = mean.view(-1, self.horizon, self.d_action)
            log_std = log_std.view(-1, self.horizon, self.d_action)
        else:
            mean = out.view(-1, self.horizon, self.d_action)

        # Ensure network output is in a reasonable range
        if not self.action_lows is None and not self.action_highs is None:
            mean = torch.tanh(mean)
            mean = mean * (self.action_highs - self.action_lows)

        # Update the shifted optimizee STD
        if self.learn_rollout_std:
            std = log_std.exp().view(-1, self.horizon, self.d_action)
            std = torch.clamp(shifted_std*std, STD_MIN, STD_MAX)

        # Optionally set the optimizee mean as a residual on the shift-forward mean
        if self.is_residual:
            mean = mean + shifted_mean

        # Return the shifted mean and STD
        if self.learn_rollout_std:
            return mean, std
        else:
            return mean

    def forward(self,
                x: Tensor,
                costs: Optional[Tensor] = None,
                mean: Optional[Tensor] = None,
                std: Optional[Tensor] = None,
                cond: Optional[Tensor] = None,
                run_critic: bool = True,
                **kwargs) -> Dict[str, Any]:
        x = utils.to_tensor(x, self.tensor_args)

        # Handle the mean input
        if not mean is None:
            self.mean = mean
        shifted_mean, old_mean = self.process_mean(x, self.mean)
        self.mean = shifted_mean

        # Handle the STD input
        if self.learn_rollout_std and not self.mppi_mode:
            if not std is None:
                self.std = std
            shifted_std, old_std = self.process_std(x, self.std)
            self.std = shifted_std
        else:
            old_std = self.process_std(x, self.std, shift=False)
            self.std = old_std

        # Use learned shift model
        if not self.shift_model is None and not self.mppi_mode:
            if self.learn_rollout_std:
                self.mean, self.std = self.run_shift_model(x, old_mean, self.mean, old_std, self.std, cond)
            else:
                self.mean = self.run_shift_model(x, old_mean,self.mean, old_std, self.std, cond)

        # Run the specified # of iterations of optimization
        for iter in range(self.n_iters):

            # Ensure we reoptimize if running multiple iterations
            if iter > 0:
                costs = None
                self.rollout_task.update_params(self.params_stacked)

            # Run rollouts to compute costs
            if costs is None:
                costs, samples = self.run_rollouts(x)
            else:
                samples = None

            # Run the actor
            mean_dist, std_dist = self.run_actor(x,
                                                 costs=costs,
                                                 mean=self.mean,
                                                 std=self.std,
                                                 cond=cond)

            # Sample the mean
            if self.use_mean:
                horizon = mean_dist.loc
            else:
                horizon = mean_dist.rsample()
            horizon = horizon.clamp(self.action_lows, self.action_highs)
            self.mean = horizon

            # Sample the STD
            if self.learn_rollout_std and not self.mppi_mode:
                if self.use_mean:
                    self.std = std_dist.loc
                else:
                    self.std = std_dist.rsample()
                self.std = self.std.clamp(STD_MIN, STD_MAX)

        # Set the action to be the first in the horizon
        action = horizon[:, 0]

        # Run the critic
        if self.mppi_mode or not run_critic:
            value = torch.zeros(x.shape[0])
        else:
            value = self.run_critic(x,
                                    costs=costs,
                                    mean=old_mean,
                                    std=old_std,
                                    cond=cond)

        return dict(
            action=action,
            horizon=horizon,
            costs=costs,
            mean=self.mean,
            old_mean=old_mean,
            std=self.std,
            old_std=old_std,
            samples=samples,
            mean_dist=mean_dist,
            std_dist=std_dist,
            value=value
        )
dmpo/controllers/dmpo_policy.py
jisacks/dmpo
import torch
from torch.utils.data import Dataset
from torch.utils.data.sampler import SubsetRandomSampler, Sampler, Sequence
import numpy as np

from .. import utils

from torch import Tensor
from torch.nn import Module
from typing import Dict, Any, Optional, List, Union, Tuple, Callable, Iterator

class SubsetSampler(Sampler[int]):
    r"""Samples elements randomly from a given list of indices, without replacement.

    Args:
        indices (sequence): a sequence of indices
        generator (Generator): Generator used in sampling.
    """
    indices: Sequence[int]

    def __init__(self, indices: Sequence[int], generator=None) -> None:
        self.indices = indices
        self.generator = generator

    def __iter__(self) -> Iterator[int]:
        for i in self.indices:
            yield self.indices[i]

    def __len__(self) -> int:
        return len(self.indices)

class DatasetBuffer(Dataset):
    """
    Dataset buffer class for PPO
    """
    def __init__(self, discount: float=1, seq_length: int=1, stride: int=1, gae_lambda: Optional[float]=None) -> None:
        """
        :param discount: Discount factor
        """
        super().__init__()
        self.discount = discount
        self.seq_length = seq_length
        self.stride = stride
        self.gae_lambda = gae_lambda

        self.clear()

    def clear(self) -> None:
        self.states = []
        self.actions = []
        self.horizons = []
        self.costs = []
        self.conds = []
        self.rewards = []
        self.values = []
        self.stds = []
        self.returns = []
        self.advantages = []
        self.means = []
        self.time_steps = []

    def push(self, trajectories: List[Dict[str, Any]]) -> None:
        # Iterate through the trajectories
        for traj in trajectories:
            self.states.append(traj['states'])
            self.actions.append(traj['actions'])
            self.costs.append(traj['costs'])
            self.rewards.append(traj['rewards'])
            self.values.append(traj['values'])
            self.conds.append(traj['conds'])
            self.horizons.append(traj['horizons'])
            self.means.append(traj['means'])
            self.stds.append(traj['stds'])
            self.time_steps.append(traj['time_steps'])

    def __len__(self) -> int:
        return len(self.states)

    def compute_returns_and_advantages(self) -> None:
        if self.gae_lambda is None:
            for traj_idx in range(len(self)):
                rewards = self.rewards[traj_idx]
                values = self.values[traj_idx]

                R = values[-1]
                returns = []
                for reward in reversed(rewards):
                    R = self.discount*R + reward
                    returns.insert(0, R)
                returns = torch.stack(returns)
                self.returns.append(returns)

            returns = torch.stack(self.returns)
            values = torch.stack(self.values)
            adv = returns - values
        else:
            for traj_idx in range(len(self)):
                rewards = self.rewards[traj_idx]
                values = self.values[traj_idx]

                last_gae_lam = 0
                advantages = []
                for t in reversed(range(len(rewards))):
                    if t == len(rewards)-1:
                        next_values = values[-1]
                    else:
                        next_values = values[t+1]
                    delta = rewards[t] + self.discount*next_values - values[t]
                    last_gae_lam = delta + self.discount*self.gae_lambda*last_gae_lam
                    advantages.insert(0, last_gae_lam)
                advantages = torch.stack(advantages)
                returns = advantages + values
                self.advantages.append(advantages)
                self.returns.append(returns)
            adv = torch.stack(self.advantages)

        adv = (adv - adv.mean()) / (adv.std() + 1e-6)
        self.advantages = adv

    def split_into_subsequences(self) -> None:
        states = self.states
        actions = self.actions
        horizons = self.horizons
        costs = self.costs
        conds = self.conds
        returns = self.returns
        rewards = self.rewards
        advantages = self.advantages
        means = self.means
        stds = self.stds
        time_steps = self.time_steps

        self.clear()

        for traj_idx in range(len(states)):
            ep_length = states[traj_idx].shape[0]
            indices = np.arange(0, ep_length-self.seq_length+1, self.stride)
            n_subseqs = len(indices)

            # Split into subsequences
            self.states.extend(
                [states[traj_idx][i:i+self.seq_length] for i in indices])
            self.actions.extend(
                [actions[traj_idx][i:i+self.seq_length] for i in indices])
            self.rewards.extend(
                [rewards[traj_idx][i:i+self.seq_length] for i in indices])
            self.time_steps.extend(
                [time_steps[traj_idx][i:i+self.seq_length] for i in indices])

            # Handle returns and advantages if computed
            if len(returns) > 0:
                self.returns.extend(
                    [returns[traj_idx][i:i+self.seq_length] for i in indices])
                self.advantages.extend(
                    [advantages[traj_idx][i:i+self.seq_length] for i in indices])
            else:
                self.returns.extend([None for _ in range(n_subseqs)])
                self.advantages.extend([None for _ in range(n_subseqs)])

            # Append the full sampled plans
            if not horizons[traj_idx] is None:
                self.horizons.extend(
                    [horizons[traj_idx][i:i+self.seq_length] for i in indices])
            else:
                self.horizons.extend([None for _ in range(n_subseqs)])

            # Handle the optional means variables
            if not means[traj_idx] is None:
                self.means.extend(
                    [means[traj_idx][i:i+self.seq_length] for i in indices])
            else:
                self.means.extend([None for _ in range(n_subseqs)])

            # Handle the optional stds variables
            if not stds[traj_idx] is None:
                self.stds.extend(
                    [stds[traj_idx][i:i + self.seq_length] for i in indices])
            else:
                self.stds.extend([None for _ in range(n_subseqs)])

            # Handle the optional costs variable
            if not costs[traj_idx] is None:
                self.costs.extend(
                    [costs[traj_idx][i:i+self.seq_length] for i in indices])
            else:
                self.costs.extend([None for _ in range(n_subseqs)])

            # Handle the optional condition variable
            if not conds[traj_idx] is None:
                self.conds.extend(
                    [conds[traj_idx][i:i+self.seq_length] for i in indices])
            else:
                self.conds.extend([None for _ in range(n_subseqs)])


    def get_samplers(self):
        indices = np.arange(len(self))
        np.random.shuffle(indices)
        sampler = SubsetSampler(indices)
        return sampler

    def __getitem__(self, idx):
        return self.states[idx], \
               self.actions[idx], \
               self.rewards[idx], \
               self.horizons[idx] if not self.horizons[idx] is None else torch.empty(()), \
               self.costs[idx] if not self.costs[idx] is None else torch.empty(()), \
               self.conds[idx] if not self.conds[idx] is None else torch.empty(()), \
               self.returns[idx] if not self.returns[idx] is None else torch.empty(()), \
               self.advantages[idx] if not self.advantages[idx] is None else torch.empty(()), \
               self.means[idx] if not self.means[idx] is None else torch.empty(()), \
               self.stds[idx] if not self.stds[idx] is None else torch.empty(()), \
               self.time_steps[idx]

dmpo/dataset/dataset_buffer.py
jisacks/dmpo
import torch
import numpy as np

# Quaternion routines adapted from rowan to use autograd
def qmultiply(q1, q2):
	return np.concatenate((
		np.array([q1[0] * q2[0] - np.sum(q1[1:4] * q2[1:4])]), # w1w2
		q1[0] * q2[1:4] + q2[0] * q1[1:4] + np.cross(q1[1:4], q2[1:4])))

def qconjugate(q):
	return np.concatenate((q[0:1],-q[1:4]))

def qrotate(q, v):
	quat_v = np.concatenate((np.array([0]), v))
	return qmultiply(q, qmultiply(quat_v, qconjugate(q)))[1:]

def qexp(q):
	norm = np.linalg.norm(q[1:4])
	e = np.exp(q[0])
	result_w = e * np.cos(norm)
	if np.isclose(norm, 0):
		result_v = np.zeros(3)
	else:
		result_v = e * q[1:4] / norm * np.sin(norm)
	return np.concatenate((np.array([result_w]), result_v))

def qintegrate(q, v, dt, frame='body'):
	quat_v = np.concatenate((np.array([0]), v*dt/2))
	if frame == 'body':
		return qmultiply(q, qexp(quat_v))		
	if frame == 'world':
		return qmultiply(qexp(quat_v), q)

def qstandardize(q):
	if q[0] < 0:
		q *= -1
	return q / np.linalg.norm(q)

def qtoR(q):
	q0 = q[0]
	q1 = q[1]
	q2 = q[2]
	q3 = q[3]
     
    # First row of the rotation matrix
	r00 = 2 * (q0 * q0 + q1 * q1) - 1
	r01 = 2 * (q1 * q2 - q0 * q3)
	r02 = 2 * (q1 * q3 + q0 * q2)
     
    # Second row of the rotation matrix
	r10 = 2 * (q1 * q2 + q0 * q3)
	r11 = 2 * (q0 * q0 + q2 * q2) - 1
	r12 = 2 * (q2 * q3 - q0 * q1)
     
    # Third row of the rotation matrix
	r20 = 2 * (q1 * q3 - q0 * q2)
	r21 = 2 * (q2 * q3 + q0 * q1)
	r22 = 2 * (q0 * q0 + q3 * q3) - 1
     
    # 3x3 rotation matrix
	rot_matrix = np.array([[r00, r01, r02], \
                           [r10, r11, r12], \
                           [r20, r21, r22]])
                            
	return rot_matrix


# torch versions:
def sqaured_distance_torch(x, x_d):
	'''
	x: actual state, tensor, (N, m)
	x_d: desired, tensor, (m,)
	output: squared distance, tensor, (N,)
	'''
	return torch.einsum('ij,ij->i', x-x_d, x-x_d)

def qdistance_torch(q1, q2):
	'''
	q1: tensor, (N, 4)
	q2: tensor, (N, 4)
	output: tensor, (N,)
	distance = 1 - <q1,q2>^2
	'''
	return 1 - torch.einsum('bi,bi->b', q1, q2)**2

def qmultiply_torch(q1, q2):
	'''
	q1: tensor, (N, 4)
	q2: tensor, (N, 4)
	output: tensor, (N, 4)
	'''
	temp = torch.zeros_like(q1)
	temp[:, 0] = q1[:,0] * q2[:,0] - torch.sum(q1[:,1:] * q2[:,1:], dim=1)
	temp[:, 1:] = q1[:,0:1] * q2[:,1:] + q2[:,0:1] * q1[:,1:] + torch.cross(q1[:,1:],q2[:,1:], dim=1)
	return temp

def qconjugate_torch(q):
	temp = torch.zeros_like(q)
	if len(q.shape) == 1:
		temp[0] = q[0]
		temp[1:] = -q[1:]
	else:
		temp[:,0] = q[:,0]
		temp[:,1:] = -q[:,1:]
	return temp

def qrotate_torch(q, v):
	'''
	q: tensor, (N, 4)
	v: tensor, (N, 3)
	output: tensor, (N, 3)
	'''
	# a more efficient way suggested by Alex:
	temp = 2. * torch.cross(q[:, 1:], v)
	return v + q[:, 0:1] * temp + torch.cross(q[:, 1:], temp)

	quat_v = torch.zeros_like(q)
	quat_v[:, 1:] = v
	return qmultiply_torch(q, qmultiply_torch(quat_v, qconjugate_torch(q)))[:,1:]

def qexp_torch(q):
	'''
	q: tensor, (N, 4)
	output: tensor, (N, 4)
	'''
	norm = torch.linalg.norm(q[:,1:], dim=1)
	e = torch.exp(q[:,0])
	result_w = e * torch.cos(norm)

	N = q.shape[0]
	result_v = torch.zeros_like(q[:,1:])
	result_v[norm>0] = e.view(N,1)[norm>0] * q[norm>0,1:] / norm.view(N,1)[norm>0] * torch.sin(norm).view(N,1)[norm>0]

	return torch.cat((result_w.view(N,1), result_v), dim=1)

def qintegrate_torch(q, v, dt, frame='body'):
	'''
	q: tensor, (N, 4)
	v: tensor, (N, 3)
	output: tensor, (N, 4)
	'''
	quat_v = torch.zeros_like(q)
	quat_v[:,1:] = v * dt / 2.
	if frame == 'body':
		return qmultiply_torch(q, qexp_torch(quat_v))
	if frame == 'world':
		return qmultiply_torch(qexp_torch(quat_v), q)

def qstandardize_torch(q):
	'''
	q: tensor, (N, 4)
	output: tensor, (N, 4)
	'''
	return torch.where(q[:, 0:1] < 0, -q, q)

def qtoR_torch(q):
	'''
	q: tensor, (N, 4)
	output: rotation matrix tensor, (N, 3, 3)
	'''
	q0 = q[:, 0]
	q1 = q[:, 1]
	q2 = q[:, 2]
	q3 = q[:, 3]
	R = torch.zeros(q.shape[0], 3, 3).to(q)

    # First row of the rotation matrix
	R[:, 0, 0] = 2 * (q0 * q0 + q1 * q1) - 1
	R[:, 0, 1] = 2 * (q1 * q2 - q0 * q3)
	R[:, 0, 2] = 2 * (q1 * q3 + q0 * q2)
     
    # Second row of the rotation matrix
	R[:, 1, 0] = 2 * (q1 * q2 + q0 * q3)
	R[:, 1, 1] = 2 * (q0 * q0 + q2 * q2) - 1
	R[:, 1, 2] = 2 * (q2 * q3 - q0 * q1)
     
    # Third row of the rotation matrix
	R[:, 2, 0] = 2 * (q1 * q3 - q0 * q2)
	R[:, 2, 1] = 2 * (q2 * q3 + q0 * q1)
	R[:, 2, 2] = 2 * (q0 * q0 + q3 * q3) - 1
                            
	return R

def get_quaternion_from_euler(roll, pitch, yaw):
	qx = np.sin(roll / 2) * np.cos(pitch / 2) * np.cos(yaw / 2) - np.cos(roll / 2) * np.sin(pitch / 2) * np.sin(yaw / 2)
	qy = np.cos(roll / 2) * np.sin(pitch / 2) * np.cos(yaw / 2) + np.sin(roll / 2) * np.cos(pitch / 2) * np.sin(yaw / 2)
	qz = np.cos(roll / 2) * np.cos(pitch / 2) * np.sin(yaw / 2) - np.sin(roll / 2) * np.sin(pitch / 2) * np.cos(yaw / 2)
	qw = np.cos(roll / 2) * np.cos(pitch / 2) * np.cos(yaw / 2) + np.sin(roll / 2) * np.sin(pitch / 2) * np.sin(yaw / 2)
	return np.stack([qx, qy, qz, qw], axis=-1)
dmpo/envs/math_utils.py
jisacks/dmpo
import meshcat
import meshcat.geometry as g
import meshcat.transformations as tf
import rowan
import time
from copy import deepcopy
import yaml
from math import ceil

from .quadrotor_param import QuadrotorParam
from .math_utils import *

from torch import Tensor
from typing import Dict, Any, Optional, List, Tuple

class QuadrotorEnv():
    """
    Quadrotor simulation environment
    """
    def __init__(self,
                 config: str,
                 num_envs: int=1,
                 use_omega: bool=False,
                 action_is_mf: bool=True,
                 convert_mf_to_omega: bool=False,
                 use_delay_model: bool=False,
                 delay_coeff: float=0.2,
                 randomize_mass: bool = False,
                 mass_range: List[float]=[0.7, 1.3],
                 randomize_delay_coeff: bool = False,
                 delay_range: List[float]=[0.2, 0.6],
                 force_pert: bool=False,
                 force_range: List[float]=[-3.5, 3.5],
                 force_is_z: bool=False,
                 ou_theta: float=0.15,
                 ou_sigma: float=0.20,
                 use_obs_noise: bool=False,
                 tensor_args: Dict[str, Any]={'device': 'cpu', 'dtype': torch.float32}):
        """
        :param config: YAML configuration file, which will be parsed to form a QuadrotorParam object
        :param num_envs: # of parallel environments to simulation
        :param use_omega: use the omega controller, which converts desired (thrust, omega) to motor forces
        :param action_is_mf: specified that the action space is motor forces
        :param convert_mf_to_omega: converts a motor force command to desired (thrust, omega) to model Crazyflie
        :param use_delay_model: use the delay model to translate desired (thrust, omega) into actual thrust and omega
        :param delay_coeff: coefficient of delay model
        :param randomize_mass: use domain randomization for mass
        :param mass_range: scaling factors for mass for domain randomization
        :param randomize_delay_coeff: use domain randomization for delay coefficient
        :param delay_range: range of delay coefficients to use for domain randomization
        :param force_pert: use random force perturbations
        :param force_range: range of force perturbations
        :param force_is_z: force perturbations can also be in Z direction (will just be XY if false)
        :param ou_theta: OU process theta parameter for changing force perturbation over time
        :param ou_sigma: OU process sigma parameter for changing force perturbation over time
        :param use_obs_noise: corrupt state with observation noise
        :param tensor_args: PyTorch tensor arguments
        """
        self.num_envs = num_envs
        self.tensor_args = tensor_args

        # Load in the configuration
        self.config = yaml.load(open(config), yaml.FullLoader)
        self.param = QuadrotorParam(self.config)

        # Action space parameters
        self.use_omega = use_omega
        self.action_is_mf = action_is_mf
        self.convert_mf_to_omega = convert_mf_to_omega

        # Delay model parameters
        self.use_delay_model = use_delay_model
        self.true_delay_coeff = delay_coeff

        # Domain randomization parameters
        self.mass_range = mass_range
        self.randomize_mass = randomize_mass
        self.delay_range = delay_range
        self.randomize_delay_coeff = randomize_delay_coeff
        self.force_pert = force_pert
        self.force_range = force_range
        self.force_is_z = force_is_z
        self.use_obs_noise = use_obs_noise
        self.ou_theta = ou_theta
        self.ou_sigma = ou_sigma

        # Init timing
        self.times = self.param.sim_times
        self.time_step = 0
        self.avg_dt = self.times[1] - self.times[0]

        # Init system state
        self.init_state = torch.tensor(self.config['initial_state'], **self.tensor_args)

        # Control bounds in motor force space
        self.a_min = torch.tensor(self.param.a_min, **self.tensor_args)
        self.a_max = torch.tensor(self.param.a_max, **self.tensor_args)

        # Get action bounds for controller (may be motor force space if specified)
        if (not self.action_is_mf and not self.convert_mf_to_omega) or use_omega:
            self.action_lows = torch.tensor([0., -10, -10, -10], **self.tensor_args)
            #self.action_highs = torch.tensor([self.a_max[0]*4, 12, 12, 12], **self.tensor_args)
            self.action_highs = torch.tensor([0.7848, 10, 10, 10], **self.tensor_args)
        else:
            self.action_lows = self.a_min
            self.action_highs = self.a_max

        # Initial conditions
        self.s_min = torch.tensor(self.param.s_min, **self.tensor_args)
        self.s_max = -self.s_min
        self.rpy_limit = torch.tensor(self.param.rpy_limit, **self.tensor_args)
        self.limits = torch.tensor(self.param.limits, **self.tensor_args)

        # Constants
        self.d_state = 13
        self.d_obs = 13
        self.d_action = 4

        self.mass = torch.tensor([self.param.mass], **self.tensor_args).repeat(self.num_envs)
        self.g = self.param.g
        self.inv_mass = 1 / self.mass

        self.d = self.param.d
        self.rho = self.param.rho
        self.Cs = self.param.Cs
        self.Ct = self.param.Ct
        self.k1 = self.param.k1
        self.k2 = self.param.k2

        self.B0 = torch.tensor(self.param.B0, **self.tensor_args)
        self.B0_inv = torch.linalg.inv(self.B0)

        self.J = torch.tensor(self.param.J, **self.tensor_args)
        if self.J.shape == (3, 3):
            self.J = torch.as_tensor(self.J, **self.tensor_args)
            self.inv_J = torch.linalg.inv(self.J)
        else:
            self.J = torch.diag(torch.as_tensor(self.J, **self.tensor_args))
            self.inv_J = torch.linalg.inv(self.J)

        # Controller gains
        self.omega_gain = self.config['omega_gain']

        # Plotting stuff
        self.states_name = [
            'Position X [m]',
            'Position Y [m]',
            'Position Z [m]',
            'Velocity X [m/s]',
            'Velocity Y [m/s]',
            'Velocity Z [m/s]',
            'qw',
            'qx',
            'qy',
            'qz',
            'Angular Velocity X [rad/s]',
            'Angular Velocity Y [rad/s]',
            'Angular Velocity Z [rad/s]']

        self.deduced_state_names = [
            'Roll [deg]',
            'Pitch [deg]',
            'Yaw [deg]',
        ]

        self.actions_name = [
            'Motor Force 1 [N]',
            'Motor Force 2 [N]',
            'Motor Force 3 [N]',
            'Motor Force 4 [N]']

        # Reward function coefficients
        # ref: row 8, Table 3, USC sim-to-real paper
        self.alpha_p = self.param.alpha_p
        self.alpha_z = self.param.alpha_z
        self.alpha_w = self.param.alpha_w
        self.alpha_a = self.param.alpha_a
        self.alpha_R = self.param.alpha_R
        self.alpha_v = self.param.alpha_v
        self.alpha_yaw = self.param.alpha_yaw
        self.alpha_pitch = self.param.alpha_pitch
        self.alpha_u_delta = self.param.alpha_u_delta
        self.alpha_u_thrust = self.param.alpha_u_thrust
        self.alpha_u_omega = self.param.alpha_u_omega

    def get_env_state(self) -> List[Tensor]:
        return self.states

    def set_env_state(self, state: List[Tensor]) -> None:
        self.states = state

    def set_param(self, param_dict: Dict[str, Any]) -> None:
        for k, v in param_dict.items():
            setattr(self, k, v)

    def reset(self) -> Tensor:
        # Determine the initial state of the quadrotor
        if self.init_state is None:
            self.states = torch.zeros(self.d_state)

            # Position and velocity
            limits = self.limits
            self.states[0:6] = torch.rand(6) * (2 * limits[0:6]) - limits[0:6]

            # Rotation
            rpy = np.radians(np.random.uniform(-self.rpy_limit, self.rpy_limit, 3))
            q = rowan.from_euler(rpy[0], rpy[1], rpy[2], 'xyz')
            self.states[6:10] = torch.tensor(q, **self.tensor_args)

            # Angular velocity
            self.states[10:13] = torch.rand(3) * (2 * limits[10:13]) - limits[10:13]
        else:
            self.states = self.init_state

        # Reset state and action variables
        self.states = self.states.unsqueeze(0).repeat(self.num_envs, 1)
        self.time_step = 0
        self.actions = torch.zeros((self.num_envs, self.d_action), **self.tensor_args)
        self.prev_actions = torch.zeros((self.num_envs, self.d_action), **self.tensor_args)

        # Randomize mass
        if self.randomize_mass:
            mass_scale = torch.rand(self.num_envs, **self.tensor_args)*(self.mass_range[1] - self.mass_range[0]) + self.mass_range[0]
            self.mass = torch.tensor([self.param.mass], **self.tensor_args).repeat(self.num_envs)
            self.mass = self.mass * mass_scale
            self.inv_mass = 1 / self.mass
        else:
            self.mass = torch.tensor([self.param.mass], **self.tensor_args).repeat(self.num_envs)
            self.inv_mass = 1 / self.mass

        # Randomize delay coeff:
        if self.randomize_delay_coeff:
            self.delay_coeff = torch.rand(self.num_envs, **self.tensor_args) * (self.delay_range[1] - self.delay_range[0]) + self.delay_range[0]
        else:
            self.delay_coeff = torch.tensor([self.true_delay_coeff], **self.tensor_args).repeat(self.num_envs)

        if self.force_pert:
            self.force_dist = torch.rand((self.num_envs, 3), **self.tensor_args) * (self.force_range[1] - self.force_range[0]) + self.force_range[0]
            if not self.force_is_z:
                self.force_dist[:, -1] = 0
        else:
            self.force_dist = torch.zeros((self.num_envs, 3), **self.tensor_args)

        # Get the reference trajectories
        self.ref_trajectory, self.ref_dts, self.ref_pos = self.param.get_reference(self.num_envs)
        self.ref_trajectory = self.ref_trajectory.to(**self.tensor_args)

        return self.states

    def get_env_obs(self) -> List[Tensor]:
        if self.use_obs_noise:
            noise = torch.randn((self.num_envs, self.d_state), **self.tensor_args)
            noise = noise * torch.as_tensor(self.param.noise_measurement_std, **self.tensor_args)
            noisystate = self.states + noise
            noisystate[:, 6:10] /= torch.norm(noisystate[:, 6:10], dim=-1, keepdim=True)
            return noisystate
        else:
            return self.states

    def f(self, s: Tensor, a: Tensor) -> Tensor:
        num_envs = s.shape[0]
        dsdt = torch.zeros(num_envs, 13).to(**self.tensor_args)
        v = s[:, 3:6]  # velocity (N, 3)
        q = s[:, 6:10]  # quaternion (N, 4)
        omega = s[:, 10:]  # angular velocity (N, 3)

        if self.action_is_mf and not self.convert_mf_to_omega:
            # If action space is motor forces and we did not convert to omega space, then compute wrench
            eta = a @ self.B0.T # output wrench (N, 4)
        else:
            # Otherwise, our action is (thrust, omega)
            eta = a

        f_u = torch.zeros(num_envs, 3).to(**self.tensor_args)
        f_u[:, 2] = eta[:, 0]  # total thrust (N, 3)
        tau_u = eta[:, 1:]  # torque (N, 3) or desired omega

        # dynamics
        # \dot{p} = v
        dsdt[:, :3] = v  # <- implies velocity and position in same frame

        # Apply the force perturbation
        if self.force_pert:
            dsdt[:, 3:6] += self.force_dist

        # mv = mg + R f_u  	# <- implies f_u in body frame, p, v in world frame
        dsdt[:, 5] -= self.g
        dsdt[:, 3:6] += qrotate_torch(q, f_u) / self.mass[:, None]

        # \dot{R} = R S(w)
        # see https://rowan.readthedocs.io/en/latest/package-calculus.html
        qnew = qintegrate_torch(q, omega, self.avg_dt, frame='body')
        qnew = qstandardize_torch(qnew)

        # transform qnew to a "delta q" that works with the usual euler integration
        dsdt[:, 6:10] = (qnew - q) / self.avg_dt

        if self.action_is_mf and not self.convert_mf_to_omega:
            # Compute omega from torques
            # J\dot{w} = Jw x w + tau_u
            Jomega = omega @ self.J.T
            dsdt[:, 10:] = torch.cross(Jomega, omega) + tau_u
            dsdt[:, 10:] = dsdt[:, 10:] @ self.inv_J.T
        else:
            # Set updated omega to be the control command
            dsdt[:, 10:] = (tau_u - omega) / self.avg_dt

        # Adding noise
        dsdt[:, 3:6] += torch.normal(mean=0,
                                     std=self.param.noise_process_std[0],
                                     size=(self.num_envs, 3),
                                     **self.tensor_args)
        dsdt[:, 10:] += torch.normal(mean=0,
                                     std=self.param.noise_process_std[1],
                                     size=(self.num_envs, 3),
                                     **self.tensor_args)

        return dsdt

    def next_state(self, s: Tensor, a: Tensor) -> Tensor:
        new_s = s + self.avg_dt * self.f(s, a)
        return new_s

    def get_cost(self, a: Tensor) -> Tensor:
        state_ref = self.ref_trajectory[:, :, self.time_step]
        p_des = state_ref[:, 0:3]
        v_des = state_ref[:, 3:6]
        w_des = state_ref[:, 10:]
        q_des = state_ref[:, 6:10]

        # Position tracking error
        if self.alpha_p > 0:
            ep = torch.linalg.norm(self.states[:, 0:3] - p_des, dim=1)
        else:
            ep = 0.

        # Additional cost on Z tracking error
        if self.alpha_z > 0:
            ez = torch.linalg.norm(self.states[:, 2:3] - p_des[2:3], dim=-1)
        else:
            ez = 0.

        # Velocity tracking error
        if self.alpha_v > 0:
            ev = torch.linalg.norm(self.states[:, 3:6] - v_des, dim=1)
        else:
            ev = 0.

        # Angular velocity tracking error
        if self.alpha_w > 0:
            ew = torch.linalg.norm(self.states[:, 10:] - w_des, dim=1)
        else:
            ew = 0.

        # Orientation tracking error
        if self.alpha_R > 0:
            eR = qdistance_torch(self.states[:, 6:10], q_des)
        else:
            eR = 0.

        # Control cost
        if self.alpha_a > 0:
            ea = torch.linalg.norm(a, dim=1)
        else:
            ea = 0.

        # Yaw tracking error
        if self.alpha_yaw > 0:
            qe = qmultiply_torch(qconjugate_torch(q_des), self.states[:, 6:10])
            Re = qtoR_torch(qe)
            eyaw = torch.atan2(Re[:, 1, 0], Re[:, 0, 0]) ** 2
        else:
            eyaw = 0.

        # Pitch tracking error
        if self.alpha_pitch > 0:
            qe = qmultiply_torch(qconjugate_torch(q_des), self.states[:, 6:10])
            Re = qtoR_torch(qe)
            epitch = (torch.asin(Re[:,2,0].clip(-1, 1)))**2
        else:
            epitch = 0

        # Penalize control changes between time steps
        if self.alpha_u_delta > 0:
            edelta = torch.norm(a - self.prev_actions, dim=1)
        else:
            edelta = 0

        # Separate thrust control cost (use only if control space includes thrust)
        if self.alpha_u_thrust > 0:
            ethrust = torch.norm(a[:, :1], dim=1)
        else:
            ethrust = 0

        # Separate omega control cost (use only if control space includes omega)
        if self.alpha_u_omega > 0:
            eomega = torch.norm(a[:, 1:], dim=1)
        else:
            eomega = 0

        cost = (self.alpha_p * ep
                + self.alpha_z * ez
                + self.alpha_v * ev
                + self.alpha_w * ew
                + self.alpha_a * ea
                + self.alpha_yaw * eyaw
                + self.alpha_R * eR
                + self.alpha_pitch * epitch
                + self.alpha_u_delta * edelta
                + self.alpha_u_thrust * ethrust
                + self.alpha_u_omega * eomega) * self.avg_dt
        return cost

    def get_env_infos(self) -> Dict[str, Any]:
        done = (self.time_step + 1) >= len(self.times)
        dones = [done for _ in range(self.num_envs)]
        return dict(dones=dones, done=done)

    def omega_controller(self, a: Tensor) -> Tensor:
        '''
        Converts desired (thrust, omega) to motor forces
        '''
        T_d = a[:, 0]
        omega_d = a[:, 1:]

        omega = self.states[:, 10:13]
        omega_e = omega_d - omega

        torque = self.omega_gain * omega_e  # tensor, (3,)
        torque = torch.mm(self.J, torque.T).T
        torque -= torch.cross(torch.mm(self.J, omega.T).T, omega)

        wrench = torch.cat((T_d.view(self.num_envs, 1), torque), dim=1)  # tensor, (N, 4)
        motorForce = torch.mm(self.B0_inv, wrench.T).T
        motorForce = torch.clip(motorForce, self.a_min, self.a_max)
        return motorForce

    def convert_motor_forces(self, a: Tensor) -> Tensor:
        '''
        Converts motor forces to desired (thrust, omega)
        '''
        eta = a @ self.B0.T
        T_d = eta[:, :1]
        tau_u = eta[:, 1:]

        omega = self.states[:, 10:]
        Jomega = omega @ self.J.T
        d_omega = torch.cross(Jomega, omega) + tau_u
        d_omega = d_omega @ self.inv_J.T
        omega_d = omega + d_omega*self.avg_dt

        new_a = torch.cat((T_d, omega_d), dim=-1)
        return new_a

    def step(self, a: Tensor) -> Tuple[Any, ...]:
        a = a.to(**self.tensor_args)

        if not self.action_is_mf and not self.convert_mf_to_omega:
            # Ensure control bounds if using desired (thrust, omega) as control space
            a = torch.clip(a, self.action_lows, self.action_highs)
        elif self.convert_mf_to_omega:
            # Convert motor forces to desired (thrust, omega) command
            a = torch.clip(a, self.a_min, self.a_max)
            a = self.convert_motor_forces(a)
        elif self.action_is_mf:
            # Ensure control bounds if using motor forces
            a = torch.clip(a, self.a_min, self.a_max)

        # Apply omega controller to convert motor forces to desired (thrust, omega)
        if self.use_omega:
            a = self.omega_controller(a)
        cmd = a.clone()

        # Apply delay model on controls (should only be used for desired thrust, omega action space)
        if self.use_delay_model:
            self.actions = self.actions + self.delay_coeff[:, None]*(a - self.actions)
            a = self.actions

        # Compute the state transitions
        new_states = self.next_state(self.states, a)
        self.states = new_states

        # Increment force perturbation with OU process
        if self.force_pert:
            d_force = -self.ou_theta * self.force_dist
            d_force += torch.randn(d_force.shape, **self.tensor_args) * self.ou_sigma
            self.force_dist += d_force * self.avg_dt
            self.force_dist = self.force_dist.clamp(self.force_range[0], self.force_range[1])
            if not self.force_is_z:
                self.force_dist[:, -1] = 0

        reward = -self.get_cost(cmd)
        info = self.get_env_infos()
        self.time_step += 1
        self.prev_actions = cmd
        return self.states, reward, info['done'], info

    def render(self,
               env_idx: int=0,
               samples: Optional[Tensor]=None,
               mean: Optional[Tensor]=None,
               state: Optional[Tensor]=None) -> Any:
        return None

    def get_param_dict(self) -> List[Dict[str, Any]]:
        """
        Return a dictionary with parameters to be sent to MPC controller
        """
        param_dict = [dict(t=self.times[self.time_step if self.time_step < len(self.times)-1 else -1],
                           actions=self.actions[idx],
                           ref_dts=self.ref_dts[idx] if not self.ref_dts is None else None,
                           ref_pos=self.ref_pos[idx] if not self.ref_pos is None else None)
                      for idx in range(self.num_envs)]
        return param_dict

    def get_env_description(self) -> List[Tensor]:
        """
        Get a Tensor containing the reference trajectory used to condition policy and value function
        """
        T = 32
        stride = 4
        dim = T//stride

        ref_traj = self.ref_trajectory[:, :, self.time_step:self.time_step+T:stride]
        ref_traj = torch.cat((ref_traj[:, :3], ref_traj[:, 6:10]), dim=1)

        if ref_traj.shape[2] < dim:
            diff = T//stride - ref_traj.shape[2]
            end_step = ref_traj[:, :, -1:].repeat(1, 1, diff)
            ref_traj = torch.cat((ref_traj, end_step), dim=-1)

        cond = ref_traj.reshape(self.num_envs, -1)
        #cond = torch.cat((cond, self.mass[:, None], self.delay_coeff[:, None], self.force_dist), dim=-1)
        info = [cond[i] for i in range(self.num_envs)]
        return info

    def evaluate_success(self, trajectories: List[Dict[str, Any]]) -> Dict[str, Any]:
        num_traj = len(trajectories)
        successes = []
        total_costs = []

        for idx, traj in enumerate(trajectories):
            costs = -traj['rewards']

            total_cost = costs.sum()
            success = True

            successes.append(success)
            total_costs.append(total_cost)
        successes = torch.tensor(successes)
        total_costs = torch.tensor(total_costs)

        success_percentage = torch.sum(successes) / num_traj * 100.
        ret_dict = dict(
            successes=successes,
            total_costs=total_costs,
            success_percentage=success_percentage
        )
        return ret_dict

    def visualize(self, states, ref_traj, dt):
        # Create a new visualizer
        vis = meshcat.Visualizer()
        vis.open()

        vis["/Cameras/default"].set_transform(
            tf.translation_matrix([0, 0, 0]).dot(
                tf.euler_matrix(0, np.radians(-30), -np.pi / 2)))

        vis["/Cameras/default/rotated/<object>"].set_transform(
            tf.translation_matrix([1, 0, 0]))

        vis['/Background'].set_property('top_color', [0, 0, 0])
        vis['/Background'].set_property('bottom_color', [0, 0, 0])

        vis["Quadrotor"].set_object(g.StlMeshGeometry.from_file('./crazyflie2.stl'))

        vertices = np.array([[0, 0.5], [0, 0], [0, 0]]).astype(np.float32)
        vis["lines_segments"].set_object(g.Line(g.PointsGeometry(vertices),
                                                g.MeshBasicMaterial(color=0xff0000, linewidth=100.)))

        vis['ref'].set_object(g.Line(g.PointsGeometry(ref_traj.numpy()[:, :3].T),
                                     g.LineBasicMaterial(color=0xff99ff, linewidth=100.)))

        while True:
            for state in states:
                vis["Quadrotor"].set_transform(
                    tf.translation_matrix([state[0], state[1], state[2]]).dot(
                        tf.quaternion_matrix(state[6:10])))

                vis["lines_segments"].set_transform(
                    tf.translation_matrix([state[0], state[1], state[2]]).dot(
                        tf.quaternion_matrix(state[6:10])))

                time.sleep(dt)
dmpo/envs/quadrotor.py
jisacks/dmpo
import numpy as np
import torch

from torch import Tensor
from numpy.typing import NDArray
from typing import Dict, Any, Optional, List, Union, Tuple

class QuadrotorParam():
	"""
	Class containing parameters for quadrotor model
	"""
	def __init__(self,
				 config: Dict[str, Any],
				 MPPI: bool=False):
		"""
		:param config: dictionary containing configuration
		:param MPPI: indicates if we should use the MPPI controller parameters (which may be different)
		"""
		self.env_name = 'Quadrotor-Crazyflie'

		# Quadrotor parameters
		self.sim_t0 = 0
		self.sim_tf = config['sim_tf']
		self.sim_dt = config['sim_dt']
		self.sim_times = np.arange(self.sim_t0, self.sim_tf, self.sim_dt)
		if MPPI:
			self.sim_dt = config['sim_dt_MPPI']
			self.sim_times = np.arange(self.sim_t0, self.sim_tf, self.sim_dt)

		# Control limits [N] for motor forces
		self.a_min = np.array(config.get('a_min', [0., 0, 0, 0]))
		self.a_max = np.array(config.get('a_max', [12., 12., 12., 12])) / 1000 * 9.81 # g->N

		# Crazyflie 2.0 quadrotor.py
		self.mass = config.get('mass', 0.034) #kg
		self.J = np.array(config.get('J', [16.571710, 16.655602, 29.261652])) * 1e-6
		self.d = 0.047

		# Side force model parameters for wind perturbations
		if config['Vwind'] == 0:
			self.wind = False
			self.Vwind = None
		else:
			self.wind = True
			self.Vwind = np.array(config['Vwind']) # velocity of wind in world frame
		self.Ct = 2.87e-3
		self.Cs = 2.31e-5
		self.k1 = 1.425
		self.k2 = 3.126
		self.rho = 1.225 # air density (in SI units)

		# Note: we assume here that our control is forces
		arm_length = 0.046 # m
		arm = 0.707106781 * arm_length
		t2t = 0.006 # thrust-to-torque ratio
		self.t2t = t2t
		self.B0 = np.array([
			[1, 1, 1, 1],
			[-arm, -arm, arm, arm],
			[-arm, arm, arm, -arm],
			[-t2t, t2t, -t2t, t2t]
			])
		self.g = 9.81 # not signed

		# Exploration parameters: state boundary and initial state sampling range
		self.s_min = np.array( \
			[-8, -8, -8, \
			  -5, -5, -5, \
			  -1.001, -1.001, -1.001, -1.001,
			  -20, -20, -20])
		self.rpy_limit = np.array([5, 5, 5])
		self.limits = np.array([0.1,0.1,0.1,0.1,0.1,0.1,0,0,0,0,0,0,0])

		# Measurement noise
		self.noise_measurement_std = np.zeros(13)
		self.noise_measurement_std[:3] = 0.005
		self.noise_measurement_std[3:6] = 0.005
		self.noise_measurement_std[6:10] = 0.01
		self.noise_measurement_std[10:] = 0.01

		# Process noise
		self.noise_process_std = config.get('noise_process_std', [0.3, 2.])

		# Reference trajectory parameters
		self.ref_type = config['traj']
		self.max_dist = config.get('max_dist', [1., 1., 0.])
		self.min_dt = config.get('min_dt', 0.6)
		self.max_dt = config.get('max_dt', 1.5)

		# Cost function parameters
		self.alpha_p = config['alpha_p']
		self.alpha_z = config['alpha_z']
		self.alpha_w = config['alpha_w']
		self.alpha_a = config['alpha_a']
		self.alpha_R = config['alpha_R']
		self.alpha_v = config['alpha_v']
		self.alpha_yaw = config['alpha_yaw']
		self.alpha_pitch = config['alpha_pitch']
		self.alpha_u_delta = config['alpha_u_delta']
		self.alpha_u_thrust = config['alpha_u_thrust']
		self.alpha_u_omega = config['alpha_u_omega']

	def get_reference(self,
					  num_envs: int,
					  dts: Optional[NDArray]=None,
					  pos: Optional[NDArray]=None) -> Tuple[Tensor, NDArray, NDArray]:
		"""
		:param num_envs: # of reference trajectories to generate (one per environment)
		:param dts: delta in time between waypoints (randomly generated if not specified)
		:param pos: waypoints for zig-zag (randomly generated if not specified)
		"""

		self.ref_trajectory = np.zeros((num_envs, 13, len(self.sim_times)))
		self.ref_trajectory[:, 6, :] = 1.

		if self.ref_type == 'zig-zag-yaw':
			if dts is None:
				dts = np.random.uniform(self.min_dt, self.max_dt, size=(num_envs, len(self.sim_times), 1))
				dts = dts.repeat(3, axis=2)
			d_times = dts.cumsum(axis=1)
			if pos is None:
				pos = np.random.uniform(0, np.array(self.max_dist), size=(num_envs, len(self.sim_times)//2, 3))

			for env_idx in range(num_envs):
				for p_idx in range(3):
					for step, time in enumerate(self.sim_times[1:]):
						ref_idx = np.searchsorted(d_times[env_idx, :, p_idx], time)
						sign = 1 if np.ceil(ref_idx/2 - ref_idx//2) == 0 else -1
						ref_pos = sign * pos[env_idx, ref_idx, p_idx]
						prev_ref_pos = -1 * sign * pos[env_idx, ref_idx-1, p_idx] if ref_idx > 0 else 0
						ref_time = d_times[env_idx, ref_idx, p_idx]
						prev_ref_time = d_times[env_idx, ref_idx-1, p_idx] if ref_idx > 0 else 0
						cur_pos = self.ref_trajectory[env_idx, p_idx, step]

						delta = ref_pos - prev_ref_pos
						if delta != 0:
							delta = delta/(ref_time-prev_ref_time)*self.sim_dt
						self.ref_trajectory[env_idx, p_idx, step+1] = cur_pos + delta
						self.ref_trajectory[env_idx, 6, step+1] = 1 if sign == 1 else 0
						self.ref_trajectory[env_idx, 9, step+1] = 1 if sign == -1 else 0
		elif self.ref_type == 'zig-zag':
			if dts is None:
				dts = np.random.uniform(self.min_dt, self.max_dt, size=(num_envs, len(self.sim_times), 1))
				dts = dts.repeat(3, axis=2)
			d_times = dts.cumsum(axis=1)
			if pos is None:
				pos = np.random.uniform(0, np.array(self.max_dist), size=(num_envs, len(self.sim_times)//2, 3))

			for env_idx in range(num_envs):
				for p_idx in range(3):
					for step, time in enumerate(self.sim_times[1:]):
						ref_idx = np.searchsorted(d_times[env_idx, :, p_idx], time)
						sign = 1 if np.ceil(ref_idx/2 - ref_idx//2) == 0 else -1
						ref_pos = sign * pos[env_idx, ref_idx, p_idx]
						prev_ref_pos = -1 * sign * pos[env_idx, ref_idx-1, p_idx] if ref_idx > 0 else 0
						ref_time = d_times[env_idx, ref_idx, p_idx]
						prev_ref_time = d_times[env_idx, ref_idx-1, p_idx] if ref_idx > 0 else 0
						cur_pos = self.ref_trajectory[env_idx, p_idx, step]

						delta = ref_pos - prev_ref_pos
						if delta != 0:
							delta = delta/(ref_time-prev_ref_time)*self.sim_dt
						self.ref_trajectory[env_idx, p_idx, step+1] = cur_pos + delta
		else:
			raise ValueError('Invalid reference trajectory type specified.')

		return torch.tensor(self.ref_trajectory), dts, pos



dmpo/envs/quadrotor_param.py
jisacks/dmpo
import torch
from pytorch_lightning.loggers import TensorBoardLogger

from .experiment_utils import create_model, create_optim
from .. import utils

from torch import Tensor
from torch.nn import Module
from typing import Dict, Any, Optional, List, Union, Tuple, Callable

class Experiment():
    """
    Base class for all experiments
    """
    def __init__(self,
                 model_config: Dict[str, Any],
                 optim_config: Dict[str, Any],
                 n_epochs: int=1,
                 log_folder: str='logs',
                 exp_name: str='experiment',
                 dtype: str='float',
                 device: str='cuda'):
        """
        :param model_config: Configuration for model
        :param optim_config: Configuration for optimizer
        :param n_epochs: # of epochs to run
        :param n_gpus: # of GPUs to utilize
        :param log_folder: Folder to use for logging
        :param exp_name: Name of experiment
        :param dtype: PyTorch datatype to use
        :param device: PyTorch device to use
        :param val_every: # of epochs per which we run the validation step
        :param max_grad_norm: Maximum gradient norm clip
        :param lr_scheduler_config: Configuration for learning rate scheduler (optional)
        """
        self.n_epochs = n_epochs

        # Set tensor args
        self.dtype = utils.TorchDtype[dtype]
        self.device = torch.device(device) if torch.cuda.is_available() else torch.device('cpu')
        self.tensor_args = {'device':self.device, 'dtype':self.dtype}

        # Create the model
        self.model = create_model(tensor_args=self.tensor_args, **model_config)
        self.model.to(**self.tensor_args)

        # Create the logger
        self.logger = TensorBoardLogger(log_folder, exp_name)

        # Get the optimizers and learning rate schedulers
        self.optim, self.optim_args = create_optim(optim_config)

dmpo/experiment/experiment.py
jisacks/dmpo
import torch
import torch.nn as nn

from ..envs.quadrotor import QuadrotorEnv
from ..mpc.task.base_rollout_task import BaseRolloutTask
from ..mpc.task.quadrotor_rollout_task import QuadrotorRolloutTask
from ..controllers.dmpo_policy import DMPOPolicy
from .. import utils

from torch.nn import Module
from typing import Dict, Any, Optional, List, Tuple, Callable

ModelTypes = ['dmpo_policy']
EnvTypes = ['quadrotor']

def create_model(model_type: str,
                 tensor_args: Dict[str, Any] = {'device': 'cpu', 'dtype': torch.float},
                 **kwargs: Dict[str, Any]) -> Module:
    if model_type in ModelTypes:
        if model_type == 'dmpo_policy':
            return DMPOPolicy(tensor_args=tensor_args, **kwargs)
    else:
        raise ValueError('Invalid model type {} specified.'.format(model_type))

def create_env(env_name: str,
               env_config_file: Optional[str]=None,
               tensor_args: Dict[str, Any]={'device': 'cpu', 'dtype': torch.float},
               **kwargs: Dict[str, Any]):
    if env_name in EnvTypes:
        # Load in the environment configuration file
        if not env_config_file is None:
            env_params = utils.load_yaml(env_config_file)

            # Allow for kwargs to overwrite configuration file
            for k, v in env_params.items():
                if not k in kwargs.keys():
                    kwargs[k] = v

        # Handle each environment separately
        if env_name == 'quadrotor':
            return QuadrotorEnv(tensor_args=tensor_args, **kwargs)
    else:
        raise ValueError('Specified environment type {} is unsupported.'.format(env_name))

def create_task(env_name: str,
                num_envs: int,
                task_config_file: str,
                tensor_args: Dict[str, Any] = {'device': 'cpu', 'dtype': torch.float},
                **kwargs: Dict[str, Any]) -> BaseRolloutTask:
    if env_name in EnvTypes:
        # Load in the task configuration file
        task_params = utils.load_yaml(task_config_file)

        # Handle each environment separately
        if env_name == 'quadrotor':
            return QuadrotorRolloutTask(exp_params=task_params,
                                        num_envs=num_envs,
                                        tensor_args=tensor_args,
                                        **kwargs)
    else:
        raise ValueError('Specified environment type {} is unsupported.'.format(env_name))

def create_optim(config: Dict[str, Any]) -> Tuple[Callable, Dict[str, Any]]:
    optim_type = config.get('optim', 'Adam')
    optim_args = config.get('optim_args', {})

    if optim_type == 'Adam':
        optim = torch.optim.Adam
    elif optim_type == 'SGD':
        optim = torch.optim.SGD
    elif optim_type == 'Adagrad':
        optim = torch.optim.Adagrad
    elif optim_type == 'RMSprop':
        optim = torch.optim.RMSprop
    elif optim_type == 'RAdam':
        optim = torch.optim.RAdam
    else:
        raise ValueError('Specified optimizer type {} unsupported.'.format(optim_type))

    return optim, optim_args
dmpo/experiment/experiment_utils.py
jisacks/dmpo
import torch
from torch.utils.data import DataLoader
import pytorch_lightning as pl
import numpy as np

from .experiment import Experiment
from .ppo_rollout import ppo_rollout
from ..dataset.dataset_buffer import DatasetBuffer
from .experiment_utils import create_task, create_env, create_optim
from .ppo_trainer import PPOTrainer
from .. import utils

from torch import Tensor
from torch.nn import Module
from typing import Dict, Any, Optional, List, Union, Tuple, Callable

class PPOExperiment(Experiment):
    """
    Experiment class for running PPO
    """
    def __init__(self,
                 env_name: str,
                 env_config: Dict[str, Any],
                 model_config: Dict[str, Any],
                 actor_optim_config: Dict[str, Any],
                 critic_optim_config: Dict[str, Any],
                 task_config: Optional[Dict[str, Any]]=None,
                 batch_size: int=1,
                 seed: int=0,
                 n_iters: int=1,
                 use_condition: bool=False,
                 train_episode_len: int=1,
                 val_episode_len: int=1,
                 train_episodes: int=1,
                 val_episodes: int=1,
                 n_val_envs: Optional[int]=None,
                 break_if_done: bool=False,
                 dynamic_env: bool=False,
                 num_workers: int=0,
                 model_file: Optional[str]=None,
                 n_epochs: int=1,
                 n_gpus: int=1,
                 log_folder: str='logs',
                 exp_name: str='experiment',
                 dtype: str='float',
                 device: str='cuda',
                 env_device: str='cpu',
                 val_every: int=1,
                 max_grad_norm: float=1,
                 n_pretrain_steps: Optional[int]=None,
                 n_pretrain_epochs: int=1,
                 dataset_config: Dict[str, Any]={},
                 trainer_config: Dict[str, Any]={},
                 train_env_config: Dict[str, Any]={},
                 val_env_config: Dict[str, Any]={}) -> None:
        """
        :param env_name: Environment to load
        :param env_config: Configuration for environment
        :param model_config: Configuration for actor\critic model
        :param actor_optim_config: Configuration for actor optimizer
        :param critic_optim_config: Configuration for critic optimizer
        :param task_config: Configuration for the MPC task
        :param batch_size: Batch size
        :param seed: Base seed for rollouts
        :param n_iters: # of iterations of rollout-train-test
        :param use_condition: Handle logic for forming and sending conditioning variable to policy and critic
        :param train_episode_len: Max length of training episodes
        :param val_episode_len: Max length of validation episodes
        :param train_episodes: # of train episodes per iteration
        :param val_episodes: # of validation episodes per iteration
        :param n_val_envs: # of validation environments (can be different than # of training envs)
        :param break_if_done: Break from rollout if all environments have finished
        :param dynamic_env: Environment is dynamic (requires additional processing for MPC)
        :param num_workers: # of threads to use for loading training data
        :param model_file: Optional pre-trained model to load
        :param n_epochs: # of epochs to train for each iteration
        :param n_gpus: # of GPU devices to use
        :param log_folder: Folder to use for logging and checkpoints
        :param exp_name: Name of experiment for logging
        :param dtype: PyTorch data type
        :param device: PyTorch device
        :param val_every: # of epochs per which we run on validation data
        :param max_grad_norm: Maximum norm used for gradient clipping
        :param n_pretrain_steps: # of iterations to pretrain critic
        :param n_pretrain_epochs: # of epochs per iteration for pretraining critic
        :param dataset_config: Configuration for dataset buffer
        :param trainer_config: Configuration for the trainer
        :param train_env_config: Additional environment configuration for training rollouts
        :param val_env_config: Additional environment configuration for validation rollouts
        """
        super().__init__(model_config=model_config,
                         optim_config=actor_optim_config,
                         n_epochs=n_epochs,
                         log_folder=log_folder,
                         exp_name=exp_name,
                         dtype=dtype,
                         device=device)

        self.batch_size = batch_size
        self.seed = seed
        self.n_iters = n_iters
        self.use_condition = use_condition
        self.train_episode_len = train_episode_len
        self.val_episode_len = val_episode_len
        self.train_episodes = train_episodes
        self.val_episodes = val_episodes
        self.val_every = val_every
        self.n_val_envs = n_val_envs
        self.break_if_done = break_if_done
        self.dynamic_env = dynamic_env
        self.num_workers = num_workers
        self.n_pretrain_steps = n_pretrain_steps
        self.n_pretrain_epochs = n_pretrain_epochs
        self.model_config = model_config
        self.task_config = task_config
        self.train_env_config = train_env_config
        self.val_env_config = val_env_config

        self.best_percent = None
        self.best_cost = None

        # Optionally load the model
        if not model_file is None:
            model_dict = torch.load(model_file, map_location='cpu')
            state_dict = model_dict['state_dict']
            self.model.load_state_dict(state_dict)
            self.model.to(**self.tensor_args)

        if hasattr(self.model, 'horizon'):
            self.horizon = self.model.horizon
        self.d_action = self.model.d_action

        # Get the critic optimizer
        self.critic_optim, self.critic_optim_args = create_optim(critic_optim_config)

        # Create the environment
        self.env_tensor_args = {'device':env_device, 'dtype':self.dtype}
        self.env = create_env(env_name, tensor_args=self.env_tensor_args, **env_config)

        self.model.action_lows = utils.to_tensor(self.env.action_lows, self.tensor_args)
        self.model.action_highs = utils.to_tensor(self.env.action_highs, self.tensor_args)

        # Create the task
        if not task_config is None:
            self.task = create_task(env_name=env_name,
                                    num_envs=self.env.num_envs,
                                    tensor_args=self.tensor_args,
                                    **task_config)
            self.model.set_task(self.task)
        else:
            self.task = None

        # Create the dataset buffer
        self.dataset = DatasetBuffer(**dataset_config)

        # Create the PPO model trainer
        self.model_trainer = PPOTrainer(model=self.model,
                                        actor_optim=self.optim,
                                        actor_optim_args=self.optim_args,
                                        critic_optim=self.critic_optim,
                                        critic_optim_args=self.critic_optim_args,
                                        max_grad_norm=max_grad_norm,
                                        **trainer_config)

        # Create the PyTorch Lightning Trainer instance
        self.trainer = pl.Trainer(devices=n_gpus,
                                  accelerator='gpu',
                                  logger=self.logger,
                                  max_epochs=self.n_epochs,
                                  check_val_every_n_epoch=val_every,
                                  num_sanity_val_steps=0)

    def run_rollouts(self, itr: int=0, is_train: bool=False):
        if is_train:
            # Set the training configuration
            n_episodes = self.train_episodes
            ep_length = self.train_episode_len
            self.env.set_param(self.train_env_config)

            # Use a different random seed for each training trial
            seed = self.seed + 123*itr
        else:
            # Set the validation configuration
            n_episodes = self.val_episodes
            ep_length = self.val_episode_len
            self.env.set_param(self.val_env_config)

            # Use a fixed random seed for each validation trial
            seed = self.seed

            # Change # of environments if different for validation
            if not self.n_val_envs is None:
                n_envs = self.env.num_envs
                self.env.num_envs = self.n_val_envs

        # Validation mode uses the mean optimizer update
        if not is_train:
            use_mean = self.model.use_mean
            self.model.use_mean = True

        # Perform the rollout
        with torch.no_grad():
            trajectories = ppo_rollout(env=self.env,
                                       model=self.model,
                                       n_episodes=n_episodes,
                                       ep_length=ep_length,
                                       base_seed=seed,
                                       break_if_done=self.break_if_done,
                                       use_condition=self.use_condition,
                                       dynamic_env=self.dynamic_env,
                                       use_tqdm=True,
                                       tensor_args=self.tensor_args)

        # Reset parameters if modified during validation mode
        if not is_train:
            self.model.use_mean = use_mean

            if not self.n_val_envs is None:
                self.env.num_envs = n_envs

        # Compute statistics
        success_dict = self.env.evaluate_success(trajectories)
        stat_dict = utils.compute_statistics(success_dict)

        # Display statistics
        success_percentage = success_dict['success_percentage']
        mean_cost = stat_dict['mean_cost']
        std_cost = stat_dict['std_cost']
        mean_success_cost = stat_dict['mean_success_cost']
        std_success_cost = stat_dict['std_success_cost']

        # Save current model
        utils.make_dir(self.logger.log_dir + '/checkpoints')

        filename = self.logger.log_dir + '/checkpoints/last.pt'
        torch.save(dict(state_dict=self.model.state_dict(),
                        model_config=self.model_config,
                        task_config=self.task_config), filename)

        # Save model if validation performance better in terms of cost
        if not is_train:
            if (self.best_cost is None) or (self.best_cost >= mean_cost):
                self.best_percent = success_percentage
                self.best_cost = mean_cost
                self.model.reset()

                utils.make_dir(self.logger.log_dir + '/checkpoints')
                filename = self.logger.log_dir + '/checkpoints/best.pt'
                torch.save(dict(state_dict=self.model.state_dict(),
                                model_config=self.model_config,
                                task_config=self.task_config),
                           filename)

        # Log validation performance
        if not is_train:
            self.logger.experiment.add_scalar('Best Success Percentage', self.best_percent, itr)
            self.logger.experiment.add_scalar('Best Mean Cost', self.best_cost, itr)

            self.logger.experiment.add_scalar('Success Percentage', success_percentage, itr)
            self.logger.experiment.add_scalar('Mean Cost', mean_cost, itr)
            self.logger.experiment.add_scalar('Mean Success Cost', mean_success_cost, itr)

        print('Success Metric (Best) = {:.2f} ({:.2f}), '
              'Mean Cost (Best) = {:.3e} +/- {:.3e} ({:.3e}), '
              'Mean Success Cost = {:.3e} +/- {:.3e}'.format(success_percentage,
                                                             self.best_percent if not self.best_percent is None else 0,
                                                             mean_cost,
                                                             std_cost,
                                                             self.best_cost if not self.best_percent is None else np.inf,
                                                             mean_success_cost,
                                                             std_success_cost))

        return trajectories

    def run(self) -> None:
        # Main loop
        for itr in range(self.n_iters):
            print('Iteration {}'.format(itr))

            # Collect training data
            trajectories = self.run_rollouts(itr=itr, is_train=True)

            # Add trajectories to dataset
            self.dataset.clear()
            self.dataset.push(trajectories)
            self.dataset.compute_returns_and_advantages()
            self.dataset.split_into_subsequences()

            # Fit the model
            sampler = self.dataset.get_samplers()
            data_loader = DataLoader(self.dataset,
                                     batch_size=self.batch_size,
                                     sampler=sampler,
                                     num_workers=self.num_workers)

            # Reset the PyTorch Lightning trainer
            self.trainer.fit_loop.epoch_progress.reset()

            # Remove the current checkpoints
            utils.remove_file(self.trainer.checkpoint_callback.best_model_path)

            last_model_path = '/'.join(self.trainer.checkpoint_callback.best_model_path.split('/')[:-1])
            last_model_path = last_model_path + '/last.ckpt'
            utils.remove_file(last_model_path)

            # Reset the PyTorch Lightning checkpoint callback
            self.trainer.checkpoint_callback.best_k_models = {}
            self.trainer.checkpoint_callback.best_model_score = None
            self.trainer.checkpoint_callback.best_model_path = ''
            self.trainer.checkpoint_callback.filename = None

            # Set up for pretraining
            if not self.n_pretrain_steps is None and itr < self.n_pretrain_steps:
                self.model_trainer.critic_only = True
                self.trainer.fit_loop.max_epochs = self.n_pretrain_epochs
            else:
                self.model_trainer.critic_only = False
                self.trainer.fit_loop.max_epochs = self.n_epochs

            # Train the model
            self.model_trainer.old_model.load_state_dict(self.model.state_dict())
            self.trainer.fit(self.model_trainer, data_loader)

            # Test current policy
            self.model.to(**self.tensor_args)
            self.model.eval()

            if itr%self.val_every == 0:
                self.run_rollouts(itr=itr, is_train=False)

dmpo/experiment/ppo_experiment.py
jisacks/dmpo
import torch
import numpy as np
from tqdm import tqdm
from time import time

from .. import utils

from torch.nn import Module
from typing import Dict, Any, Optional, List

def ppo_rollout(env,
                model: Module,
                n_episodes: int=1,
                ep_length: int=1,
                base_seed: int=0,
                break_if_done: bool=True,
                use_condition: bool=False,
                dynamic_env: bool=False,
                use_tqdm: bool=True,
                save_data: bool=True,
                run_critic=True,
                tensor_args: Dict[str, Any]={'device': 'cpu', 'dtype': torch.float32}) -> List[Dict[str, Any]]:
    """
    Rollout controller in environment
    :param env: Environment to rollout in
    :param model: Container of actor to rollout and critic
    :param n_episodes: # of episodes
    :param ep_length: Episode length
    :param base_seed: Base seed (which is modified per episode)
    :param break_if_done: Break from episode loop if done signal has been achieved
    :param use_condition: Extract conditional information from environment to give to the policy
    :param dynamic_env: Run a dynamic environment
    :param save_data: Flag to indicate if we should return all rollout data
    :param run_critic: Flag to indicate if we should run the critic
    :param tensor_args: PyTorch Tensor settings
    """
    trajectories = []

    for ep in range(n_episodes):

        # Set the seed
        episode_seed = base_seed + 12345*ep
        np.random.seed(episode_seed)
        torch.random.manual_seed(episode_seed)
        if tensor_args['device'] == 'cuda':
            torch.cuda.manual_seed(episode_seed)

        # Reset the environment
        env.reset()
        param_dict = env.get_param_dict()

        # Optional settings if using MPC under the hood
        if hasattr(model, 'update_params'):
            model.update_params(param_dict)

        if hasattr(model, 'set_seed'):
            model.set_seed(episode_seed)

        # Retrieve environment information for conditioning model
        if use_condition:
            env_info = env.get_env_description()
            cond = torch.stack(env_info, dim=0).to(**tensor_args)

            if hasattr(model, 'set_cond'):
                model.set_cond(env_info)
        else:
            cond = None

        # Reset the actor and critic
        if hasattr(model, 'reset'):
            model.reset()

        # Perform an episode
        state_list = []; cond_list = []; reward_list = []; info_list = []
        action_list = []; horizon_list = []; costs_list = []; value_list = []
        old_mean_list = []; old_std_list = []
        time_step_list = []

        pbar = tqdm(range(ep_length)) if use_tqdm else range(ep_length)
        for t in pbar:
            # Get current state
            curr_states = env.get_env_state()

            if hasattr(env, 'get_env_obs'):
                curr_obs = env.get_env_obs()
            else:
                curr_obs = curr_states

            # Run the actor and critic
            out = model(curr_obs, cond=cond, run_critic=run_critic)

            # Retrieve model outputs
            action = out['action']
            value = out['value']
            horizon = out['horizon'] if 'horizon' in out.keys() else None
            old_mean = out['old_mean'] if 'old_mean' in out.keys() else None
            old_std = out['old_std'] if 'old_std' in out.keys() else None
            costs = out['costs'] if 'costs' in out.keys() else None

            # Prepare action to be applied
            if isinstance(action, list) and isinstance(action[0], torch.Tensor):
                action = torch.stack(action).cpu()

            # Take a step in the environment
            with torch.no_grad():
                obs, reward, done, info = env.step(action)

            # Store information
            state_list.append(torch.stack(curr_states).cpu() if not isinstance(curr_states, torch.Tensor) else curr_states.cpu())
            action_list.append(action.cpu())
            reward_list.append(reward.cpu())
            info_list.append(info)

            if save_data:
                value_list.append(value.cpu())
                time_step_list.append(t)

                if not horizon is None:
                    horizon_list.append(horizon.cpu())

                if not cond is None:
                    cond_list.append(cond.cpu())

                if not costs is None:
                    costs_list.append(costs.cpu())
                else:
                    costs_list.append(None)

                if not old_mean is None:
                    old_mean_list.append(old_mean.cpu())
                else:
                    old_mean_list.append(None)

                if not old_std is None:
                    old_std_list.append(old_std.cpu())
                else:
                    old_std_list.append(None)

            # Handle a dynamic environment
            if dynamic_env:
                param_dict = env.get_param_dict()
                if hasattr(model, 'update_params'):
                    model.update_params(param_dict)

                if use_condition:
                    env_info = env.get_env_description()
                    cond = torch.stack(env_info, dim=0).to(**tensor_args)

                    if hasattr(model, 'set_cond'):
                        model.set_cond(env_info)

            # Break if done
            if break_if_done and done:
                break

        # Store trajectories
        states = utils.stack_list_tensors(state_list)
        actions = utils.stack_list_tensors(action_list)
        rewards = utils.stack_list_tensors(reward_list)
        infos = utils.transpose_dict_list(info_list)

        if save_data:
            horizons = utils.stack_list_tensors(horizon_list)
            values = utils.stack_list_tensors(value_list)
            time_steps = [torch.tensor(time_step_list).unsqueeze(1) for _ in range(env.num_envs)]

            if not cond is None:
                conds = utils.stack_list_tensors(cond_list)
            else:
                conds = None

            if not costs_list[0] is None:
                costs = utils.stack_list_tensors(costs_list)
            else:
                costs = None

            if not old_mean_list[0] is None:
                old_means = utils.stack_list_tensors(old_mean_list)
            else:
                old_means = None

            if not old_std_list[0] is None:
                old_stds = utils.stack_list_tensors(old_std_list)
            else:
                old_stds = None

        for idx in range(env.num_envs):
            if save_data:
                traj = dict(
                    states=states[idx],
                    rewards=rewards[idx],
                    infos=infos[idx],
                    actions=actions[idx],
                    horizons=horizons[idx],
                    costs=costs[idx] if not costs is None else None,
                    values=values[idx],
                    conds=conds[idx] if not conds is None else None,
                    means=old_means[idx] if old_means is not None else None,
                    stds=old_stds[idx] if not old_stds is None else None,
                    time_steps=time_steps[idx]
                )
            else:
                traj = dict(
                    states=states[idx],
                    actions=actions[idx],
                    rewards=rewards[idx],
                    infos=infos[idx]
                )
            trajectories.append(traj)
    return trajectories

dmpo/experiment/ppo_rollout.py
jisacks/dmpo
import torch
import torch.distributions as D
from copy import deepcopy

import pytorch_lightning as pl
from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from pytorch_lightning.callbacks.callback import Callback

from torch import Tensor
from torch.nn import Module
from typing import Dict, Any, Optional, List, Union, Tuple, Callable

class PPOTrainer(pl.LightningModule):
    """
    PyTorch Lightning module which implements PPO for updating the policy
    """
    def __init__(self,
                 model: Module,
                 actor_optim: Callable,
                 actor_optim_args: Dict[str, Any],
                 critic_optim: Callable,
                 critic_optim_args: Dict[str, Any],
                 entropy_penalty: float,
                 kl_penalty: float,
                 clip_epsilon: Optional[float] = None,
                 std_clip_epsilon: Optional[float] = None,
                 max_grad_norm: Optional[float] = None,
                 model_subsets: Optional[List[Union[str, List[str]]]] = None,
                 critic_only: bool=False):
        super().__init__()

        # Actor and critic optimizers and associated arguments
        self.optims = [actor_optim, critic_optim]
        self.optims_args = [actor_optim_args, critic_optim_args]
        self.model_subsets = model_subsets

        # PPO parameters
        self.clip_epsilon = clip_epsilon
        self.std_clip_epsilon = std_clip_epsilon
        self.entropy_penalty = entropy_penalty
        self.kl_penalty = kl_penalty
        self.max_grad_norm = max_grad_norm
        self.critic_only = critic_only

        # Turn off PyTorch Lightning automatic optimization
        self.automatic_optimization = False

        # Create a copy of the model for computing CPI loss
        self.model = model
        if hasattr(self.model, 'rollout_task'):
            self.rollout_task = self.model.rollout_task
            self.model.rollout_task = None

            self.old_model = deepcopy(self.model)

            self.model.rollout_task = self.rollout_task
            self.old_model.rollout_task = self.rollout_task
        else:
            self.old_model = deepcopy(self.model)

    def training_step(self, batch: List[Any], batch_idx: int) -> Dict[str, Any]:
        # Get data from batch
        states, actions, rewards, horizons, costs, conds, returns, advantages, means, stds, time_steps = batch
        batch_size, T, H, d_action = horizons.shape

        # Get the optimizers
        actor_opt, critic_opt = self.optimizers()

        # Reset the models
        if hasattr(self.model, 'reset'):
            self.model.reset()
            self.old_model.reset()

        # Iterate over subsequence for actor update
        old_log_probs_list = []; old_std_log_probs_list = []
        log_probs_list = []; std_log_probs_list = []
        entropy_list = []; kl_div_list = []

        for t in range(T):
            state = states[:, t]
            cond = conds[:, t] if conds.ndim > 1 else None
            horizon = horizons[:, t]
            cost = costs[:, t] if costs.ndim > 1 else None
            mean = means[:, t] if means.ndim > 1 else None
            std = stds[:, t] if stds.ndim > 1 else None

            # Get old action log probabilities
            with torch.no_grad():
                self.old_model.to('cuda')
                out = self.old_model(state, cond=cond, costs=cost, mean=mean, std=std)
                old_dist = out['mean_dist']
                old_std_dist = out['std_dist'] if 'std_dist' in out.keys() else None

                old_log_probs = old_dist.log_prob(horizon)
                old_log_probs = old_log_probs.mean(dim=(2)).mean(dim=(1))
                old_log_probs_list.append(old_log_probs)

                if not old_std_dist is None:
                    old_std_log_probs = old_std_dist.log_prob(std)
                    old_std_log_probs = old_std_log_probs.mean(dim=(2)).mean(dim=(1))
                    old_std_log_probs_list.append(old_std_log_probs)

            # Get new action log probabilities
            out = self.model(state, cond=cond, costs=cost, mean=mean, std=std)
            dist = out['mean_dist']
            std_dist = out['std_dist'] if 'std_dist' in out.keys() else None

            log_probs = dist.log_prob(horizon)
            log_probs = log_probs.mean(dim=(2)).mean(dim=(1))
            log_probs_list.append(log_probs)

            if not std_dist is None:
                std_log_probs = std_dist.log_prob(std)
                std_log_probs = std_log_probs.mean(dim=(2)).mean(dim=(1))
                std_log_probs_list.append(std_log_probs)

            # Compute entropy and KL divergence
            if not dist is None:
                entropy = dist.entropy()
                kl_div = D.kl_divergence(dist, old_dist)

                entropy_list.append(entropy)
                kl_div_list.append(kl_div)

        # Stack the lists
        old_log_probs = torch.stack(old_log_probs_list, dim=1)
        log_probs = torch.stack(log_probs_list, dim=1)
        entropies = torch.stack(entropy_list, dim=1)
        kl_divs = torch.stack(kl_div_list, dim=1)

        # Compute the actor loss
        ratio = torch.exp(log_probs - old_log_probs)
        cpi_loss = ratio * advantages.squeeze(2)

        # Compute actor loss for learning STD
        if not std_dist is None:
            old_std_log_probs = torch.stack(old_std_log_probs_list, dim=1)
            std_log_probs = torch.stack(std_log_probs_list, dim=1)
            std_ratio = torch.exp(std_log_probs - old_std_log_probs)
            std_cpi_loss = std_ratio * advantages.squeeze(2)
        else:
            std_ratio = None

        if not self.clip_epsilon is None:
            clip_loss = ratio.clamp(1 - self.clip_epsilon, 1 + self.clip_epsilon) * advantages.squeeze(2)
            actor_loss = -torch.min(cpi_loss, clip_loss).mean()

            if not std_ratio is None:
                clip_loss = std_ratio.clamp(1 - self.std_clip_epsilon,
                                            1 + self.std_clip_epsilon) * advantages.squeeze(2)
                actor_loss += -torch.min(std_cpi_loss, clip_loss).mean()
        else:
            actor_loss = -cpi_loss.mean() - std_cpi_loss.mean()

        # Compute entropy penalty and KL divergence
        entropy = entropies.mean()
        kl_div = kl_divs.mean()
        total_loss = actor_loss - self.entropy_penalty * entropy + self.kl_penalty * kl_div

        # Compute ratio deviation
        ratio_dev = torch.abs(ratio - 1).mean()

        # Optimize the actor
        actor_opt.zero_grad()

        if not self.critic_only:
            self.manual_backward(total_loss)
            if not self.max_grad_norm is None:
                torch.nn.utils.clip_grad_norm_(actor_opt.param_groups[0]['params'], self.max_grad_norm)
            actor_opt.step()

        actor_opt.zero_grad()
        critic_opt.zero_grad()

        # Reset the model
        if hasattr(self.model, 'reset'):
            self.model.reset()
            self.old_model.reset()

        # Iterate over subsequence for critic update
        values_list = []
        for t in range(T):
            state = states[:, t]
            cond = conds[:, t] if conds.ndim > 1 else None
            cost = costs[:, t] if costs.ndim > 1 else None
            mean = means[:, t] if means.ndim > 1 else None
            std = stds[:, t] if stds.ndim > 1 else None

            out = self.model(state, cond=cond, costs=cost, mean=mean, std=std)
            value = out['value']
            values_list.append(value)

        # Stack the list
        values = torch.stack(values_list, dim=1)

        # Compute the critic loss
        critic_loss = 0.5 * (returns - values).pow(2).mean()

        # Optimize the critic
        critic_opt.zero_grad()
        self.manual_backward(critic_loss)
        if not self.max_grad_norm is None:
            torch.nn.utils.clip_grad_norm_(critic_opt.param_groups[0]['params'], self.max_grad_norm)
        critic_opt.step()

        actor_opt.zero_grad()
        critic_opt.zero_grad()

        # Reset the model
        if hasattr(self.model, 'reset'):
            self.model.reset()
            self.old_model.reset()

        loss_dict = dict(
            actor_loss=actor_loss.detach(),
            critic_loss=critic_loss.detach(),
            total_loss=total_loss.detach(),
            ratio=ratio.mean().detach(),
            ratio_dev=ratio_dev.detach(),
            entropy=entropy.detach(),
            kl_div=kl_div.detach()
        )

        if not std_ratio is None:
            loss_dict['std_ratio'] = std_ratio.mean().detach()

        batch_dict = {'loss': total_loss, 'log': loss_dict}
        self.log_dict(loss_dict, on_step=False, on_epoch=True, prog_bar=True, logger=True)
        return batch_dict

    def configure_optimizers(self) -> List[Dict[str, Any]]:
        optimizers = []
        for idx, optim in enumerate(self.optims):
            # Get the proper subset of parameters for each optimizer
            if not self.model_subsets is None:
                model_subset = self.model_subsets

                if isinstance(model_subset[idx], str):
                    # If a single attribute is a subset
                    params = filter(lambda x: x.requires_grad, getattr(self.model, model_subset[idx]).parameters())
                else:
                    # If we have a list of attributes for the subset
                    params = []
                    for subset in model_subset[idx]:
                        params.extend(filter(lambda x: x.requires_grad, getattr(self.model, subset).parameters()))
            else:
                # Use all parameters for the optimizer
                params = filter(lambda x: x.requires_grad, self.model.parameters())

            # Create the optimizer
            optimizer = optim(params, **self.optims_args[idx])
            optimizers.append(optimizer)
        opt_dict = [{'optimizer': optimizer} for optimizer in optimizers]
        return opt_dict
dmpo/experiment/ppo_trainer.py
jisacks/dmpo
import torch.nn as nn

from torch import Tensor
from typing import Optional, List, Union, Dict, Any

from .. import utils

class MLP(nn.Module):
    """
    Class for constructing multilayer perceptrons
    """
    def __init__(self,
                 in_size: int,
                 out_size: int,
                 hidden_size: List[int],
                 act: Union[str, nn.Module],
                 dropout_prob: Optional[float]=None,
                 init_scale: Optional[float]=None,
                 act_params: Dict[str, Any]={},
                 last_linear: bool=True) -> None:
        super().__init__()
        self.in_size = in_size
        self.out_size = out_size
        self.hidden_size = hidden_size

        # Convert activation string to Module
        if isinstance(act, str) and act != 'identity':
            act = utils.ActivationType[act]

        # Create the MLP
        net = []
        prev_dim = in_size
        if not hidden_size is None and len(hidden_size) > 0:
            for hidden in hidden_size:
                # Construct the linear layer
                layer = nn.Linear(prev_dim, hidden)
                net.append(layer)

                # Apply dropout
                if not dropout_prob is None:
                    net.append(nn.Dropout(p=dropout_prob))

                # Append activation
                if act != 'identity':
                    net.append(act(**act_params))

                prev_dim = hidden

        # Apply terminal linear layer
        if last_linear:
            layer = nn.Linear(prev_dim, out_size)
            if not init_scale is None:
                layer.weight.data.normal_(0, init_scale)
                layer.bias.data.fill_(0)
            net.append(layer)

        self.net = nn.Sequential(*net)

    def forward(self, x: Tensor) -> Tensor:
        # Automatically reshape input to be (batch_size, input_dim)
        x_shape = x.shape
        x = x.reshape(-1, x_shape[-1])
        out = self.net(x)

        # Reshape input to original shape
        out = out.reshape(*x_shape[:-1], -1)
        return out
dmpo/models/mlp.py
jisacks/dmpo
from torch.nn import Module
from .mlp import MLP

from typing import Optional

def create_net(net_type: str, in_size: Optional[int]=None, out_size: Optional[int]=None, **kwargs) -> Module:
    if net_type == 'mlp':
        net = MLP(in_size=in_size, out_size=out_size, **kwargs)
    else:
        raise ValueError('Invalid network type {} specified'.format(net_type))
    return net

dmpo/models/net_utils.py
jisacks/dmpo
import torch
import numpy as np
import yaml

from ...envs.quadrotor_param import QuadrotorParam
from ...envs.math_utils import *

from torch import Tensor
from numpy.typing import NDArray
from typing import Dict, Any, Optional, List, Union, Tuple

class QuadrotorModel():
    '''
    Quadrotor model used by MPC
    '''
    def __init__(self,
                 config: str,
                 num_envs: int=1,
                 use_omega: bool = False,
                 action_is_mf: bool = True,
                 convert_mf_to_omega: bool = False,
                 use_delay_model: bool = False,
                 delay_coeff: float = 0.4,
                 tensor_args: Dict[str, Any]={'device': 'cpu', 'dtype': torch.float32}):
        """
        :param config: YAML configuration file, which will be parsed to form a QuadrotorParam object
        :param num_envs: # of parallel environments to simulation
        :param use_omega: use the omega controller, which converts desired (thrust, omega) to motor forces
        :param action_is_mf: specified that the action space is motor forces
        :param convert_mf_to_omega: converts a motor force command to desired (thrust, omega) to model Crazyflie
        :param use_delay_model: use the delay model to translate desired (thrust, omega) into actual thrust and omega
        :param delay_coeff: coefficient of delay model
        :param tensor_args: PyTorch tensor arguments
        """
        super().__init__()
        self.num_envs = num_envs
        self.use_omega = use_omega
        self.action_is_mf = action_is_mf
        self.convert_mf_to_omega = convert_mf_to_omega
        self.use_delay_model = use_delay_model
        self.delay_coeff = delay_coeff
        self.tensor_args = tensor_args

        # Get the quadrotor configuration
        self.config = yaml.load(open(config), yaml.FullLoader)
        self.param = QuadrotorParam(self.config, MPPI=True)

        # Init timing
        self.times = self.param.sim_times
        self.time_step = 0
        self.avg_dt = self.times[1] - self.times[0]

        # Init system state
        self.init_state = torch.tensor(self.config['initial_state'], **self.tensor_args)

        # Control bounds
        self.a_min = torch.tensor(self.param.a_min, **self.tensor_args)
        self.a_max = torch.tensor(self.param.a_max, **self.tensor_args)

        if (not self.action_is_mf and not self.convert_mf_to_omega) or use_omega:
            self.action_lows = torch.tensor([0., -10, -10, -10], **self.tensor_args)
            #self.action_highs = torch.tensor([self.a_max[0]*4, 12, 12, 12], **self.tensor_args)
            self.action_highs = torch.tensor([0.7848, 10, 10, 10], **self.tensor_args)
        else:
            self.action_lows = self.a_min
            self.action_highs = self.a_max

        # Initial conditions
        self.s_min = torch.tensor(self.param.s_min, **self.tensor_args)
        self.s_max = -self.s_min
        self.rpy_limit = torch.tensor(self.param.rpy_limit, **self.tensor_args)
        self.limits = torch.tensor(self.param.limits, **self.tensor_args)

        # Constants
        self.d_state = 13
        self.d_obs = 13
        self.d_action = 4

        self.mass = self.param.mass
        self.g = self.param.g
        self.inv_mass = 1 / self.mass

        self.d = self.param.d
        self.rho = self.param.rho
        self.Cs = self.param.Cs
        self.Ct = self.param.Ct
        self.k1 = self.param.k1
        self.k2 = self.param.k2

        self.B0 = torch.tensor(self.param.B0, **self.tensor_args)
        self.B0_inv = torch.linalg.inv(self.B0)

        self.J = torch.tensor(self.param.J, **self.tensor_args)
        if self.J.shape == (3, 3):
            self.J = torch.as_tensor(self.J, **self.tensor_args)
            self.inv_J = torch.linalg.inv(self.J)
        else:
            self.J = torch.diag(torch.as_tensor(self.J, **self.tensor_args))
            self.inv_J = torch.linalg.inv(self.J)

        # Controller gains
        self.omega_gain = self.config['omega_gain']

        # Plotting stuff
        self.states_name = [
            'Position X [m]',
            'Position Y [m]',
            'Position Z [m]',
            'Velocity X [m/s]',
            'Velocity Y [m/s]',
            'Velocity Z [m/s]',
            'qw',
            'qx',
            'qy',
            'qz',
            'Angular Velocity X [rad/s]',
            'Angular Velocity Y [rad/s]',
            'Angular Velocity Z [rad/s]']

        self.deduced_state_names = [
            'Roll [deg]',
            'Pitch [deg]',
            'Yaw [deg]',
        ]

        self.actions_name = [
            'Motor Force 1 [N]',
            'Motor Force 2 [N]',
            'Motor Force 3 [N]',
            'Motor Force 4 [N]']

    def f(self, s: Tensor, a: Tensor) -> Tensor:
        num_envs, num_samples, d_state = s.shape
        dsdt = torch.zeros(num_envs*num_samples, 13).to(**self.tensor_args)
        v = s[:, :, 3:6].view(-1, 3)  # velocity (N, 3)
        q = s[:, :, 6:10].view(-1, 4)  # quaternion (N, 4)
        omega = s[:, :, 10:].view(-1, 3)  # angular velocity (N, 3)
        a = a.view(-1, 4)

        if self.action_is_mf and not self.convert_mf_to_omega:
            # If action space is motor forces and we did not convert to omega space, then compute wrench
            eta = a @ self.B0.T # output wrench (N, 4)
        else:
            # Otherwise, our action is (thrust, omega)
            eta = a

        f_u = torch.zeros(num_envs*num_samples, 3).to(**self.tensor_args)
        f_u[:, 2] = eta[:, 0]  # total thrust (N, 3)
        tau_u = eta[:, 1:]  # torque (N, 3)

        # dynamics
        # \dot{p} = v
        dsdt[:, :3] = v  # <- implies velocity and position in same frame

        # mv = mg + R f_u  	# <- implies f_u in body frame, p, v in world frame
        dsdt[:, 5] -= self.g
        dsdt[:, 3:6] += qrotate_torch(q, f_u) / self.mass

        # \dot{R} = R S(w)
        # see https://rowan.readthedocs.io/en/latest/package-calculus.html
        qnew = qintegrate_torch(q, omega, self.avg_dt, frame='body')
        qnew = qstandardize_torch(qnew)

        # transform qnew to a "delta q" that works with the usual euler integration
        dsdt[:, 6:10] = (qnew - q) / self.avg_dt

        if self.action_is_mf and not self.convert_mf_to_omega:
            # Compute omega from torques
            # J\dot{w} = Jw x w + tau_u
            Jomega = omega @ self.J.T
            dsdt[:, 10:] = torch.cross(Jomega, omega) + tau_u
            dsdt[:, 10:] = dsdt[:, 10:] @ self.inv_J.T
        else:
            # Set updated omega to be the control command
            dsdt[:, 10:] = (tau_u - omega) / self.avg_dt

        dsdt = dsdt.view(num_envs, num_samples, -1)
        return dsdt

    def step(self, s: Tensor, a: Tensor) -> Tensor:
        new_s = s + self.avg_dt * self.f(s, a)
        return new_s

    def rollout_open_loop(self, start_state: Tensor, act_seq: Tensor) -> Dict[str, Any]:
        num_particles, horizon, _ = act_seq.shape
        state_seq = torch.zeros((num_particles, horizon, self.d_state), **self.tensor_args)
        state_t = start_state.repeat((num_particles, 1))

        for t in range(horizon):
           state_t = self.step(state_t, act_seq[:,t])
           state_seq[:, t] = state_t

        trajectories = dict(
            actions = act_seq,
            state_seq = state_seq
        )

        return trajectories
    
    def get_next_state(self, curr_state, act, dt):
        pass
dmpo/mpc/model/quadrotor_model.py
jisacks/dmpo
from .rollout_base import RolloutBase
from ..model.quadrotor_model import QuadrotorModel
from ...envs.math_utils import *

import torch

from torch import Tensor
from typing import Dict, Any, Optional, List

class QuadrotorRollout(RolloutBase):
    def __init__(self,
                 dynamics_model: QuadrotorModel,
                 exp_params: Dict[str, Any],
                 num_envs: int = 1,
                 tensor_args: Dict[str, Any] = {'device': 'cpu', 'dtype': torch.float32}):
        """
        :param dynamics_model: Dynamics model used in rollout
        :param exp_params: Cost function parameters
        :param num_envs: Number of environments
        :param tensor_args: PyTorch params
        """
        super(QuadrotorRollout, self).__init__()
        self.dynamics_model = dynamics_model
        self.tensor_args = tensor_args
        self.exp_params = exp_params
        self.num_envs = num_envs

        self.use_omega = dynamics_model.use_omega
        self.action_is_mf = dynamics_model.action_is_mf
        self.convert_mf_to_omega = dynamics_model.convert_mf_to_omega
        self.use_delay_model = dynamics_model.use_delay_model

        self.param = self.dynamics_model.param
        self.alpha_p = self.param.alpha_p
        self.alpha_z = self.param.alpha_z
        self.alpha_w = self.param.alpha_w
        self.alpha_a = self.param.alpha_a
        self.alpha_R = self.param.alpha_R
        self.alpha_v = self.param.alpha_v
        self.alpha_pitch = self.param.alpha_pitch
        self.alpha_yaw = self.param.alpha_yaw
        self.t = torch.zeros((self.num_envs), **self.tensor_args)

    def cost_fn(self, state: Tensor, act: Tensor) -> Dict[str, Tensor]:
        time_step = torch.ceil(self.t / self.dynamics_model.avg_dt)
        num_envs, num_samples, H, _ = state.shape
        indices = torch.arange(H).to(**self.tensor_args) + time_step[0]
        indices = indices.clip(0, self.ref_trajectory.shape[-1]-1)
        indices = indices.to(torch.long)

        state_ref = self.ref_trajectory[:, :, indices].permute((0, 2, 1))
        p_des = state_ref[:, :, 0:3]
        v_des = state_ref[:, :, 3:6]
        w_des = state_ref[:, :, 10:]
        q_des = state_ref[:, :, 6:10]

        # Position tracking error
        if self.alpha_p > 0:
            ep = torch.linalg.norm(state[:, :, :, 0:3] - p_des[:, None], dim=-1)
        else:
            ep = 0.

        # Additional cost on Z tracking error
        if self.alpha_z > 0:
            ez = torch.linalg.norm(state[:, :, :, 2:3] - p_des[:, None, :, 2:3], dim=-1)
        else:
            ez = 0.

        # Velocity tracking error
        if self.alpha_v > 0:
            ev = torch.linalg.norm(state[:, :, :, 3:6] - v_des[:, None], dim=-1)
        else:
            ev = 0.

        # Angular velocity tracking error
        if self.alpha_w > 0:
            ew = torch.linalg.norm(state[:, :, :, 10:] - w_des[:, None], dim=-1)
        else:
            ew = 0.

        # Orientation tracking error
        if self.alpha_R > 0:
            q_des_repeated = q_des[:, None].repeat(1, num_samples, 1, 1)
            eR = qdistance_torch(state[:, :, :, 6:10].view(-1, 4), q_des_repeated.view(-1, 4))
            eR = eR.view(num_envs, num_samples, H)
        else:
            eR = 0.

        # Control cost
        if self.alpha_a > 0:
            ea = torch.linalg.norm(act, dim=-1)
        else:
            ea = 0.

        # Yaw tracking error
        if self.alpha_yaw > 0:
            q_des_repeated = q_des[:, None, :, :].repeat(1, num_samples, 1, 1).view(-1, 4)
            qe = qmultiply_torch(qconjugate_torch(q_des_repeated), state[:, :, :, 6:10].view(-1, 4))
            Re = qtoR_torch(qe)
            eyaw = torch.atan2(Re[:, 1, 0], Re[:, 0, 0]) ** 2
            eyaw = eyaw.view(num_envs, num_samples, H)
        else:
            eyaw = 0.

        # Pitch tracking error
        if self.alpha_pitch > 0:
            q_des_repeated = q_des[:, None, :, :].repeat(1, num_samples, 1, 1).view(-1, 4)
            qe = qmultiply_torch(qconjugate_torch(q_des_repeated), state[:, :, :, 6:10].view(-1, 4))
            Re = qtoR_torch(qe)
            epitch = (torch.asin(Re[:,2,0].clip(-1, 1)))**2
            epitch = epitch.view(num_envs, num_samples, H)
        else:
            epitch = 0.

        cost = (self.alpha_p * ep
                + self.alpha_z * ez
                + self.alpha_v * ev
                + self.alpha_w * ew
                + self.alpha_a * ea
                + self.alpha_yaw * eyaw
                + self.alpha_R * eR
                + self.alpha_pitch * epitch) * self.dynamics_model.avg_dt

        return dict(
            cost=cost,
            ep=self.alpha_p * ep * self.dynamics_model.avg_dt,
            ez=self.alpha_z * ez * self.dynamics_model.avg_dt,
            ev=self.alpha_v * ev * self.dynamics_model.avg_dt,
            ew=self.alpha_w * ew * self.dynamics_model.avg_dt,
            ea=self.alpha_a * ea * self.dynamics_model.avg_dt,
            eyaw=self.alpha_yaw * eyaw * self.dynamics_model.avg_dt,
            eR=self.alpha_R * eR * self.dynamics_model.avg_dt,
            epitch=self.alpha_pitch * epitch * self.dynamics_model.avg_dt
        )

    def omega_controller(self, s: Tensor, a: Tensor) -> Tensor:
        # Converts desired (thrust, omega) to motor forces
        num_envs, num_samples, _ = a.shape

        T_d = a[:, :, 0].reshape(-1)
        omega_d = a[:, :, 1:].reshape(-1, 3)
        omega = s[:, :, 10:13].reshape(-1, 3)
        omega_e = omega_d - omega

        torque = self.dynamics_model.omega_gain * omega_e  # tensor, (3,)
        torque = torch.mm(self.dynamics_model.J, torque.T).T
        torque -= torch.cross(torch.mm(self.dynamics_model.J, omega.T).T, omega)

        wrench = torch.cat((T_d.view(-1, 1), torque), dim=1)  # tensor, (N, 4)
        motorForce = torch.mm(self.dynamics_model.B0_inv, wrench.T).T
        motorForce = torch.clip(motorForce, self.dynamics_model.a_min, self.dynamics_model.a_max)
        motorForce = motorForce.view(num_envs, num_samples, -1)
        return motorForce

    def convert_motor_forces(self, s: Tensor, a: Tensor) -> Tensor:
        '''
        Converts motor forces to desired (thrust, omega)
        '''
        num_envs, num_particles, _ = a.shape
        eta = a.view(num_envs*num_particles, -1) @ self.dynamics_model.B0.T
        T_d = eta[:, :1]
        tau_u = eta[:, 1:]

        omega = s[:, :, 10:].view(num_envs*num_particles, -1)
        Jomega = omega @ self.dynamics_model.J.T
        d_omega = torch.cross(Jomega, omega) + tau_u
        d_omega = d_omega @ self.dynamics_model.inv_J.T
        omega_d = omega + d_omega*self.dynamics_model.avg_dt

        new_a = torch.cat((T_d, omega_d), dim=-1)
        new_a = new_a.view(num_envs, num_particles, -1)
        return new_a

    def rollout_fn(self, start_state: Tensor, act_seq: Tensor) -> Dict[str, Any]:
        actions = act_seq
        if actions.ndim == 3:
            actions = actions.unsqueeze(0)

        num_envs, num_particles, horizon, _ = actions.shape
        state_seq = torch.zeros((num_envs, num_particles, horizon, start_state.shape[-1]), **self.tensor_args)
        states = start_state.unsqueeze(1).repeat((1, num_particles, 1))

        # Apply delay model if action space is not motor forces
        if not self.action_is_mf and self.use_delay_model:
            if self.actions is None:
                self.actions = self.env_actions.unsqueeze(2).repeat(1, num_particles, horizon, 1)
            action = torch.clip(actions, self.dynamics_model.action_lows, self.dynamics_model.action_highs)
            self.actions = self.actions + self.dynamics_model.delay_coeff * (actions - self.actions)

        for t in range(horizon):
            if not self.action_is_mf and not self.convert_mf_to_omega:
                # Ensure control bounds if using desired (thrust, omega) as control space
                action = actions[:, :, t]
                action = torch.clip(action, self.dynamics_model.action_lows, self.dynamics_model.action_highs)
            elif self.convert_mf_to_omega:
                # Convert motor forces to desired (thrust, omega) command
                action = actions[:, :, t]
                action = torch.clip(action, self.dynamics_model.a_min, self.dynamics_model.a_max)
                action = self.convert_motor_forces(states, action)
            elif self.action_is_mf:
                # Ensure control bounds if using motor forces
                action = actions[:, :, t]
                action = torch.clip(action, self.dynamics_model.a_min, self.dynamics_model.a_max)

            if self.use_omega:
                # Apply omega controller to convert motor forces to desired (thrust, omega)
                action = self.omega_controller(states, actions[:, :, t])
            elif self.use_delay_model:
                # Get action from the delayed actions
                action = self.actions[:, :, t]

            states = self.dynamics_model.step(states, action)
            state_seq[:, :, t] = states

        cost_seq = self.cost_fn(state_seq, act_seq)

        trajectories = dict(
            actions=act_seq,
            costs=cost_seq,
            rollout_time=0.0,
            state_seq=state_seq
        )

        return trajectories

    def __call__(self, start_state: Tensor, act_seq: Tensor) -> Dict[str, Any]:
        return self.rollout_fn(start_state, act_seq)

    def update_params(self, t, actions, ref_dts, ref_pos):
        self.t = torch.tensor(t, **self.tensor_args)

        # Reset stored actions if at initial time point
        self.env_actions = torch.stack(actions).to(**self.tensor_args)[:, None]
        if np.any(t) == 0:
           self.actions = None

        # Get the reference trajectory
        kwargs = {}
        if not ref_dts is None:
            kwargs['dts'] = np.stack(ref_dts)
        if not ref_pos is None:
            kwargs['pos'] = np.stack(ref_pos)

        self.num_envs = len(ref_pos)
        self.ref_trajectory, _, _ = self.param.get_reference(self.num_envs, **kwargs)
        self.ref_trajectory = self.ref_trajectory.to(**self.tensor_args)
        return True
dmpo/mpc/rollout/quadrotor_rollout.py
jisacks/dmpo
class RolloutBase:
    def __init__(self):
        pass

    def cost_fn(self, state, act):
        pass

    def rollout_fn(self, state, act):
        pass

    def current_cost(self, current_state):
        pass

    def update_params(self):
        pass

dmpo/mpc/rollout/rollout_base.py
jisacks/dmpo
import torch

from torch import Tensor
from typing import Optional, Union, Dict, Any, Callable

class RolloutGenerator():
    """
    Class which handles rolling out dynamics models
    """
    def __init__(self,
                 rollout_fn: Optional[Callable[[Tensor, Tensor], dict]]=None,
                 tensor_args: Dict[str, Any]={'device': 'cpu', 'dtype': torch.float32}) -> None:
        self._rollout_fn = rollout_fn
        self.tensor_args = tensor_args

    @property
    def rollout_fn(self):
        return self._rollout_fn

    @rollout_fn.setter
    def rollout_fn(self, fn: Callable[[Tensor, Tensor], dict]):
        self._rollout_fn = fn

    def run_rollouts(self, state: Tensor, act_seq: Tensor) -> Dict[str, Any]:
        state = state.to(**self.tensor_args)
        act_seq = act_seq.to(**self.tensor_args)

        trajectories = self._rollout_fn(state, act_seq)
        return trajectories

    def update_params(self, kwargs: Dict[str, Any]) -> bool:
        return self.rollout_fn.update_params(**kwargs)
dmpo/mpc/rollout/rollout_generator.py
jisacks/dmpo
import torch

from ..rollout.rollout_generator import RolloutGenerator

from torch import Tensor
from typing import Optional, Dict, Any, List

class BaseRolloutTask():
    """
    Base class for defining MPC tasks
    """
    def __init__(self,
                 exp_params: Dict[str, Any],
                 num_envs: Optional[int]=None,
                 tensor_args: Dict[str, Any]={'device': "cpu", 'dtype': torch.float32}) -> None:
        """
        :param exp_params: Rollout parameters
        :param num_envs: # of environments (threads)
        :param tensor_args: PyTorch params
        """
        self.tensor_args = tensor_args
        self.num_envs = num_envs
        self.running: bool = False
        self.top_idx: Optional[List[Tensor]] = None
        self.top_values: Optional[List[Tensor]] = None
        self.top_trajs: Optional[List[Tensor]] = None

        self.rollout_fn = self.init_rollout(exp_params)
        self.init_aux(exp_params, num_envs)

    def init_aux(self, exp_params: Dict[str, Any], num_envs: int):
        pass

    def get_rollout_fn(self, exp_params: Dict[str, Any]):
        raise NotImplementedError('Function get_rollout_fn has not implemented.')

    def init_rollout(self, exp_params):
        rollout_fn = self.get_rollout_fn(exp_params)
        return RolloutGenerator(rollout_fn=rollout_fn, tensor_args=self.tensor_args)

    def run_rollouts(self, states: Tensor, act_seqs: Tensor) -> Dict[str, Any]:
        trajectories = self.rollout_fn.run_rollouts(states, act_seqs)
        return trajectories

    def update_params(self, kwargs: Dict[str, Any]) -> bool:
        self.rollout_fn.update_params(kwargs)
        return True

dmpo/mpc/task/base_rollout_task.py
jisacks/dmpo
import torch

from .base_rollout_task import BaseRolloutTask
from ..model.quadrotor_model import QuadrotorModel
from ..rollout.quadrotor_rollout import QuadrotorRollout

from typing import Optional, Dict, Any

class QuadrotorRolloutTask(BaseRolloutTask):
    def __init__(self,
                 exp_params: Dict[str, Any],
                 num_envs: Optional[int]=None,
                 tensor_args: Dict[str, Any]={'device': "cpu", 'dtype': torch.float32},
                 **kwargs) -> None:
        super().__init__(
            exp_params=exp_params,
            num_envs=num_envs,
            tensor_args=tensor_args
        )

    def get_rollout_fn(self, exp_params: Dict[str, Any]):
        dynamics_params = exp_params['model']
        self.dynamics_model = QuadrotorModel(tensor_args=self.tensor_args, num_envs=self.num_envs, **dynamics_params)
        rollout_fn = QuadrotorRollout(dynamics_model=self.dynamics_model,
                                      tensor_args=self.tensor_args,
                                      num_envs=self.num_envs,
                                      exp_params=exp_params)
        return rollout_fn

dmpo/mpc/task/quadrotor_rollout_task.py
jisacks/dmpo
import torch
import torch.nn as nn
import numpy as np
import os
import copy
import yaml
import ghalton

from torch import Tensor
from numpy.typing import NDArray
from typing import Optional, Tuple, Dict, Any, Union, List

ActivationType = {'relu': nn.ReLU,
                  'sigmoid': nn.Sigmoid,
                  'tanh': nn.Tanh,
                  'elu': nn.ELU,
                  'leakyrelu': nn.LeakyReLU,
                  'prelu': nn.PReLU,
                  'relu6': nn.ReLU6,
                  'selu': nn.SELU,
                  'celu': nn.CELU,
                  'gelu': nn.GELU,
                  'silu': nn.SiLU,
                  'mish': nn.Mish,
                  'softplus': nn.Softplus,
                  'tanhshrink': nn.Tanhshrink
                  }
TorchDtype = {'float': torch.float, 'double': torch.double, 'long': torch.long}

def compute_statistics(success_dict: Dict[str, Any]) -> Dict[str, Any]:
    successes = success_dict['successes']
    total_costs = success_dict['total_costs']

    mean_success_cost = torch.mean(total_costs[successes])
    std_success_cost = torch.std(total_costs[successes])
    mean_fail_cost = torch.mean(total_costs[~successes])
    std_fail_cost = torch.std(total_costs[~successes])
    mean_cost = torch.mean(total_costs)
    std_cost = torch.std(total_costs)
    ret_dict = dict(
        mean_success_cost=mean_success_cost,
        std_success_cost=std_success_cost,
        mean_fail_cost=mean_fail_cost,
        std_fail_cost=std_fail_cost,
        mean_cost=mean_cost,
        std_cost=std_cost
    )
    return ret_dict

def make_dir(path: str):
    """
    Create a directory if the given path does not exist
    """
    if not os.path.exists(path):
        os.makedirs(path)

def remove_file(path: str):
    if os.path.exists(path):
        os.remove(path)

def load_yaml(filename: str) -> Dict[str, Any]:
    with open(filename) as file:
        params = yaml.load(file, Loader=yaml.FullLoader)
    return params

def to_tensor(input: Union[List[Union[Tensor, NDArray, float]], Tensor, NDArray, float],
              tensor_args={'device':'cpu', 'dtype':torch.float32}) -> Tensor:
    if isinstance(input, list):
        if isinstance(input[0], list):
            return torch.tensor(input, **tensor_args)
        elif isinstance(input[0], np.ndarray):
            return torch.stack([to_tensor(x) for x in input])
        elif isinstance(input[0], torch.Tensor):
            return torch.stack(input).to(**tensor_args)
        elif isinstance(input[0], float):
            return torch.tensor(input, **tensor_args)
        else:
            raise ValueError('Invalid input to convert to tensor.')
    elif isinstance(input, torch.Tensor):
        return input.to(**tensor_args)
    elif isinstance(input, np.ndarray):
        return torch.from_numpy(input).to(**tensor_args)
    elif isinstance(input, float):
        return torch.tensor(input, **tensor_args)
    else:
        raise ValueError('Invalid input to convert to tensor.')

def as_list(x: Union[List[Any], Any], n: int =1) -> List[Any]:
    """
    Return a variable as a list
    :param x: Variable to be copied n times in a list (returns x if already a list)
    :param n: # of times to copy the variable
    """
    if isinstance(x, list):
        return x
    else:
        return [x for _ in range(n)]

def stack_list_tensors(input: List[Tensor]) -> List[Tensor]:
    """
    Converts a list (length T) of Tensors (N, ...) into a list (length N) of Tensors (T, ...)
    """
    out = []
    for idx in range(len(input[0])):
        tensor_list = torch.stack([x[idx] for x in input], dim=0)
        out.append(tensor_list)
    return out

def transpose_dict_list(input: List[Dict[str, Any]]) -> List[Dict[str, List[Any]]]:
    """
    Converts a list (length T) of dictionaries of lists\Tensors (length N) into a list (length N) of dictionaries
    consisting of lists (length T)
    """
    keys = list(input[0].keys())
    length = len(input[0][keys[0]])
    out = list({} for _ in range(length))
    for k in keys:
        y = [x[k] for x in input]
        for idx in range(length):
            v = [data[idx] if isinstance(data, list) or isinstance(data, torch.Tensor) else data for data in y]
            v = [x.detach() for x in v if isinstance(x, torch.Tensor)]
            out[idx][k] = copy.deepcopy(v)
    return out

def stack_list_list_dict_tensors(input: List[List[Dict[str, Tensor]]]) -> List[Dict[str, Tensor]]:
    """
    Converts a list (length T) of lists (length N) of dictionaries of Tensors to a list (length N) of dictionaries of
    Tensors (T, ...)
    """
    keys = list(input[0][0].keys())
    length = len(input[0])
    out = list({} for _ in range(length))
    for k in keys:
        for idx in range(length):
            out[idx][k] = []
            for t in range(len(input)):
                y = [x[k] for x in input[t]]
                out[idx][k].append(copy.deepcopy(y[idx]))
            out[idx][k] = torch.stack(out[idx][k])
    return out

def load_struct_from_dict(struct_instance, dict_instance):
    for key in dict_instance.keys():
        if (hasattr(struct_instance, key)):
            if (isinstance(dict_instance[key], dict)):
                sub_struct = load_struct_from_dict(getattr(struct_instance, key), dict_instance[key])
                setattr(struct_instance, key, sub_struct)
            else:
                setattr(struct_instance, key, dict_instance[key])
    return struct_instance

def set_if_empty(dict, key, value):
    if not key in dict.keys():
        dict[key] = value
    return dict

def generate_gaussian_halton_samples(num_samples: int,
                                     ndims: int,
                                     seed_val: int=123,
                                     device='cpu',
                                     dtype=torch.float32) -> Tensor:
    """
    Generate Halton sequence and transform to Gaussian distribution
    :param num_samples: # of samples
    :param ndims: # of independent Gaussians from which to sample
    :param seed_val: Seed for the Halton sequence generator
    :param device: PyTorch device
    :param dtype: PyTorch dtype
    """
    sequencer = ghalton.GeneralizedHalton(ndims, seed_val)
    uniform_halton_samples = torch.tensor(sequencer.get(num_samples), device=device, dtype=dtype)
    gaussian_halton_samples = torch.sqrt(torch.tensor([2.0], device=device, dtype=dtype)) \
                              * torch.erfinv(2 * uniform_halton_samples - 1)
    return gaussian_halton_samples
dmpo/utils.py
jisacks/dmpo
import yaml
import argparse
import os

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--config', type=str, help='Config file to load.')
    args = parser.parse_args()

    # Load the configuration file
    config_file = args.config
    config = yaml.load(open(config_file, 'r'), Loader=yaml.FullLoader)

    # Set the seed
    seed = config.get('seed', 0)

    # Create the experiment
    from dmpo.experiment.ppo_experiment import PPOExperiment
    exp = PPOExperiment(**config)

    import numpy as np
    import torch
    np.random.seed(seed)
    torch.manual_seed(seed)

    exp.run()
scripts/ppo_main.py
jisacks/dmpo
import torch

from dmpo.envs.quadrotor import QuadrotorEnv
from dmpo.experiment.ppo_rollout import ppo_rollout
from dmpo.experiment.experiment_utils import create_model, create_task
from dmpo import utils

if __name__ == '__main__':
    # Main script configuration
    mpc_file = '../config/mpc/quadrotor_zigzagyaw_mppi.yml'
    train_file = '../config/experiments/quadrotor_dmpo_zigzagyaw.yml'
    model_file = 'quadrotor_logs/quadrotor_dmpo_zigzagyaw/version_0/checkpoints/best.pt'
    use_gpu = True
    is_mppi = True

    # Set up PyTorch
    device = 'cuda' if use_gpu else 'cpu'
    env_tensor_args = {'device': device, 'dtype': torch.double}
    ctrl_tensor_args = {'device': device, 'dtype': torch.double}

    # Get environment
    exp_params = utils.load_yaml(mpc_file)
    rollout_params = exp_params['rollout']
    env_params = exp_params['environment']
    env_cost_params = exp_params['env_cost']
    cost_params = exp_params['cost']

    env_type = rollout_params.pop('env_type')
    num_envs = rollout_params.pop('num_envs')
    env = QuadrotorEnv(num_envs=num_envs, tensor_args=env_tensor_args, **env_params)

    # Instantiate controller from configuration
    if not is_mppi:
        model_dict = torch.load(model_file, map_location='cpu')

        state_dict = model_dict['state_dict']
        model_config = model_dict['model_config']

        train_config = utils.load_yaml(train_file)
        task_config = train_config['task_config']
    else:
        train_config = utils.load_yaml(train_file)
        model_config = train_config['model_config']
        task_config = train_config['task_config']

    if is_mppi:
        model_config['mppi_mode'] = True
        model_config['horizon'] = 32
        model_config['num_particles'] = 2048
        model_config['n_iters'] = 1
        task_config['task_config_file'] = mpc_file

    model = create_model(tensor_args=ctrl_tensor_args, **model_config)
    if not is_mppi:
        model.load_state_dict(state_dict)
    model.to(**ctrl_tensor_args)

    model.action_lows = utils.to_tensor(env.action_lows, ctrl_tensor_args)
    model.action_highs = utils.to_tensor(env.action_highs, ctrl_tensor_args)
    model.use_mean = True

    # Create the MPC task for performing rollouts
    task = create_task(env_name=env_type,
                       num_envs=num_envs,
                       tensor_args=ctrl_tensor_args,
                       **task_config)
    model.set_task(task)

    # Perform the rollouts
    with torch.no_grad():
        trajectories = ppo_rollout(env=env,
                                   model=model,
                                   tensor_args=ctrl_tensor_args,
                                   save_data=False,
                                   run_critic=False,
                                   **rollout_params)

    # Compute statistics
    success_dict = env.evaluate_success(trajectories)
    stat_dict = utils.compute_statistics(success_dict)

    success_percentage = success_dict['success_percentage']
    mean_success_cost = stat_dict['mean_cost']
    std_success_cost = stat_dict['std_cost']

    print('Success Metric = {:.2f}, Mean Cost = {:.3e} +/- {:.3e}'.format(
        success_percentage,
        mean_success_cost,
        std_success_cost,
    ))

    # Visualize trajectory
    states = trajectories[0]['states']

    mean_state_samples = None
    mppi_state_samples = None

    ref_traj = env.ref_trajectory[0].cpu().T
    env.visualize(states, ref_traj, env.avg_dt)

scripts/run_dmpo_quadrotor.py
jisacks/dmpo
# Runs the installation
from setuptools import find_packages, setup

# Avoids duplication of requirements
with open("requirements.txt") as file:
    requirements = file.read().splitlines()

setup(
    name="dmpo",
    author="Jacob Sacks",
    author_email="jsacks6@cs.washington.edu",
    description="PyTorch code for the paper Deep Model Predictive Optimization",
    url="https://github.com/jisacks/dmpo",
    install_requires=requirements,
    include_package_data=True,
    python_requires=">=3.8",
    version='1.0.0',
    packages=find_packages(),
)
setup.py
jisacks/dmpo