Linear Dynamics and Sampling Distribution Unit Tests

include/mppi/dynamics/linear/linear.cu

@@ -0,0 +1,99 @@
+/**
+ * @file linear.cu
+ * @brief Linear Dynamics
+ * @author Bogdan Vlahov
+ * @version
+ * @date 2024-08-16
+ */
+#include <mppi/dynamics/racer_dubins/racer_dubins.cuh>
+#include <string>
+#include <mppi/utils/math_utils.h>
+
+namespace mm = mppi::matrix_multiplication;
+namespace mp1 = mppi::p1;
+
+#define LINEAR_TEMPLATE template <class CLASS_T, class PARAMS_T>
+#define LINEAR_DYNAMICS LinearDynamicsImpl<CLASS_T, PARAMS_T>
+
+LINEAR_TEMPLATE LINEAR_DYNAMICS::LinearDynamicsImpl(cudaStream_t stream) : PARENT_CLASS(stream)
+{
+  this->SHARED_MEM_REQUEST_GRD_BYTES = sizeof(PARAMS_T);
+}
+
+LINEAR_TEMPLATE LINEAR_DYNAMICS::LinearDynamicsImpl(PARAMS_T& params, cudaStream_t stream)
+  : PARENT_CLASS(params, stream)
+{
+  this->SHARED_MEM_REQUEST_GRD_BYTES = sizeof(PARAMS_T);
+}
+
+LINEAR_TEMPLATE
+bool LINEAR_DYNAMICS::computeGrad(const Eigen::Ref<const state_array>& state,
+                                  const Eigen::Ref<const control_array>& control, Eigen::Ref<dfdx> A,
+                                  Eigen::Ref<dfdu> B)
+{
+  A = this->params_.getA();
+  B = this->params_.getB();
+  return true;
+}
+
+LINEAR_TEMPLATE
+LINEAR_DYNAMICS::state_array LINEAR_DYNAMICS::stateFromMap(const std::map<std::string, float>& map)
+{
+  state_array x = this->getZeroState();
+  for (int i = 0; i < this->STATE_DIM; i++)
+  {
+    std::string state_name = "x_" + std::to_string(i);
+    x(i, 0) = map.at(state_name);
+  }
+  return x;
+}
+
+LINEAR_TEMPLATE
+void LINEAR_DYNAMICS::step(Eigen::Ref<state_array> state, Eigen::Ref<state_array> next_state,
+                           Eigen::Ref<state_array> state_der, const Eigen::Ref<const control_array>& control,
+                           Eigen::Ref<output_array> output, const float t, const float dt)
+{
+  state_der = this->params_.getA() * state + this->params_.getB() * control;
+  next_state = state + state_der * dt;
+  this->stateToOutput(next_state, output);
+}
+
+LINEAR_TEMPLATE
+__device__ inline void LINEAR_DYNAMICS::step(float* state, float* next_state, float* state_der, float* control,
+                                             float* output, float* theta_s, const float t, const float dt)
+{
+  PARAMS_T* params_p = &(this->params_);
+  if (this->getGrdSharedSizeBytes() != 0)
+  {
+    params_p = (PARAMS_T*)theta_s;
+  }
+  float* A = params_p->A;
+  float* B = params_p->B;
+
+  const mp1::Parallel1Dir P_DIR = mp1::Parallel1Dir::THREAD_Y;
+  mm::gemm1<this->STATE_DIM, this->STATE_DIM, 1, P_DIR>(A, state, state_der);
+  mm::gemm1<this->STATE_DIM, this->CONTROL_DIM, 1, P_DIR>(B, control, state_der, 1.0f, 1.0f);
+  // __syncthreads();
+  int index, step;
+  mp1::getParallel1DIndex<P_DIR>(index, step);
+  for (int i = index; i < this->STATE_DIM; i += step)
+  {
+    next_state[i] = state[i] + state_der[i] * dt;
+    output[i] = next_state[i];
+  }
+}
+
+LINEAR_TEMPLATE
+__device__ void LINEAR_DYNAMICS::initializeDynamics(float* state, float* control, float* output, float* theta_s,
+                                                    float t_0, float dt)
+{
+  PARENT_CLASS::initializeDynamics(state, control, output, theta_s, t_0, dt);
+  if (this->getGrdSharedSizeBytes() != 0)
+  {
+    PARAMS_T* shared = (PARAMS_T*)theta_s;
+    *shared = this->params_;
+  }
+}
+
+#undef LINEAR_TEMPLATE
+#undef LINEAR_DYNAMICS

include/mppi/dynamics/linear/linear.cuh

@@ -0,0 +1,137 @@
+/**
+ * @file linear.cuh
+ * @brief Linear Dynamics
+ * @author Bogdan Vlahov
+ * @version
+ * @date 2024-08-16
+ */
+#pragma once
+
+#include <mppi/dynamics/dynamics.cuh>
+
+template <int STATE_DIM = 1, int CONTROL_DIM = 1>
+struct LinearDynamicsParams : public DynamicsParams
+{
+public:
+  using state_matrix = Eigen::Matrix<float, STATE_DIM, STATE_DIM>;
+  using control_matrix = Eigen::Matrix<float, STATE_DIM, CONTROL_DIM>;
+  enum class StateIndex : int
+  {
+    NUM_STATES = STATE_DIM,
+  };
+
+  enum class ControlIndex : int
+  {
+    NUM_CONTROLS = CONTROL_DIM,
+  };
+  enum class OutputIndex : int
+  {
+    NUM_OUTPUTS = STATE_DIM,
+  };
+  float A[STATE_DIM * STATE_DIM] = { 0.0f };
+  float B[STATE_DIM * CONTROL_DIM] = { 0.0f };
+
+  LinearDynamicsParams() = default;
+  ~LinearDynamicsParams() = default;
+
+  void setA(const Eigen::Ref<const state_matrix>& A_eigen)
+  {
+    memcpy(A, A_eigen.data(), sizeof(float) * STATE_DIM * STATE_DIM);
+  }
+
+  void setB(const Eigen::Ref<const control_matrix>& B_eigen)
+  {
+    memcpy(B, B_eigen.data(), sizeof(float) * STATE_DIM * CONTROL_DIM);
+  }
+
+  float* getA() const
+  {
+    return A;
+  }
+  float* getB() const
+  {
+    return B;
+  }
+
+  Eigen::Map<state_matrix> getA()
+  {
+    Eigen::Map<state_matrix> A_eigen(A);
+    return A_eigen;
+  }
+
+  Eigen::Map<control_matrix> getB()
+  {
+    Eigen::Map<control_matrix> B_eigen(B);
+    return B_eigen;
+  }
+};
+
+using namespace MPPI_internal;
+
+template <class CLASS_T, class PARAMS_T = LinearDynamicsParams<1, 1>>
+class LinearDynamicsImpl : public Dynamics<CLASS_T, PARAMS_T>
+{
+public:
+  using PARENT_CLASS = Dynamics<CLASS_T, PARAMS_T>;
+  using state_array = typename PARENT_CLASS::state_array;
+  using control_array = typename PARENT_CLASS::control_array;
+  using output_array = typename PARENT_CLASS::output_array;
+  using dfdx = typename PARENT_CLASS::dfdx;
+  using dfdu = typename PARENT_CLASS::dfdu;
+  using PARENT_CLASS::initializeDynamics;
+
+  LinearDynamicsImpl(cudaStream_t stream = nullptr);
+  LinearDynamicsImpl(PARAMS_T& params, cudaStream_t stream = nullptr);
+
+  std::string getDynamicsModelName() const override
+  {
+    return "Linear Dynamics";
+  }
+
+  bool computeGrad(const Eigen::Ref<const state_array>& state, const Eigen::Ref<const control_array>& control,
+                   Eigen::Ref<dfdx> A, Eigen::Ref<dfdu> B);
+
+  state_array stateFromMap(const std::map<std::string, float>& map);
+
+  void step(Eigen::Ref<state_array> state, Eigen::Ref<state_array> next_state, Eigen::Ref<state_array> state_der,
+            const Eigen::Ref<const control_array>& control, Eigen::Ref<output_array> output, const float t,
+            const float dt);
+
+  void setA(const Eigen::Ref<const dfdx>& A_eigen, bool synchronize = false)
+  {
+    this->params_.setA(A_eigen);
+    this->paramsToDevice(synchronize);
+  }
+
+  void setB(const Eigen::Ref<const dfdu>& B_eigen, bool synchronize = false)
+  {
+    this->params_.setB(B_eigen);
+    this->paramsToDevice(synchronize);
+  }
+
+  __device__ inline void step(float* state, float* next_state, float* state_der, float* control, float* output,
+                              float* theta_s, const float t, const float dt);
+
+  __device__ void initializeDynamics(float* state, float* control, float* output, float* theta_s, float t_0, float dt);
+};
+
+template <int STATE_DIM = 1, int CONTROL_DIM = 1>
+class LinearDynamics
+  : public LinearDynamicsImpl<LinearDynamics<STATE_DIM, CONTROL_DIM>, LinearDynamicsParams<STATE_DIM, CONTROL_DIM>>
+{
+public:
+  using PARENT_CLASS =
+      LinearDynamicsImpl<LinearDynamics<STATE_DIM, CONTROL_DIM>, LinearDynamicsParams<STATE_DIM, CONTROL_DIM>>;
+  using DYN_PARAMS_T = typename PARENT_CLASS::DYN_PARAMS_T;
+
+  LinearDynamics(cudaStream_t stream = nullptr) : PARENT_CLASS(stream)
+  {
+  }
+  LinearDynamics(DYN_PARAMS_T& params, cudaStream_t stream = nullptr) : PARENT_CLASS(params, stream)
+  {
+  }
+};
+
+#ifdef __CUDACC__
+#include "linear.cu"
+#endif

include/mppi/sampling_distributions/colored_noise/colored_noise.cu

@@ -298,7 +298,7 @@ __host__ void COLORED_NOISE::generateSamples(const int& optimization_stride, con
   const int freq_size = sample_freq.size();
 
   int smaller_index = 0;
-  const int local_control_dim = this->CONTROL_DIM; // Needed for methods which use pass by reference
+  const int local_control_dim = this->CONTROL_DIM;  // Needed for methods which use pass by reference
   Eigen::MatrixXf sample_freqs(freq_size, local_control_dim);
 
   // Adjust the weighting of each frequency by the exponents

include/mppi/sampling_distributions/gaussian/gaussian.cu

@@ -495,13 +495,8 @@ __host__ __device__ float GAUSSIAN_CLASS::computeLikelihoodRatioCost(const float
   float* control_cost_coeff = params_p->control_cost_coeff;
 
   float cost = 0.0f;
-#ifdef __CUDA_ARCH__
-  int i = threadIdx.y;
-  int step = blockDim.y;
-#else
-  int i = 0;
-  int step = 1;
-#endif
+  int i, step;
+  mppi::p1::getParallel1DIndex<mppi::p1::Parallel1Dir::THREAD_Y>(i, step);
 
   if (CONTROL_DIM % 4 == 0)
   {
@@ -584,13 +579,8 @@ __host__ __device__ float GAUSSIAN_CLASS::computeFeedbackCost(const float* __res
   float* control_cost_coeff = params_p->control_cost_coeff;
 
   float cost = 0.0f;
-#ifdef __CUDA_ARCH__
-  int i = threadIdx.y;
-  int step = blockDim.y;
-#else
-  int i = 0;
-  int step = 1;
-#endif
+  int i, step;
+  mppi::p1::getParallel1DIndex<mppi::p1::Parallel1Dir::THREAD_Y>(i, step);
 
   if (CONTROL_DIM % 4 == 0)
   {
@@ -629,7 +619,26 @@ __host__ __device__ float GAUSSIAN_CLASS::computeFeedbackCost(const float* __res
 }
 
 GAUSSIAN_TEMPLATE
-__host__ float GAUSSIAN_CLASS::computeLikelihoodRatioCost(const Eigen::Ref<const control_array>& u, const int t,
+__host__ float GAUSSIAN_CLASS::computeFeedbackCost(const Eigen::Ref<const control_array>& u_fb, const int t,
+                                                   const int distribution_idx, const float lambda, const float alpha)
+{
+  float cost = 0.0f;
+  const int distribution_i = distribution_idx >= this->params_.num_distributions ? 0 : distribution_idx;
+  float* std_dev = &(this->params_.std_dev[CONTROL_DIM * distribution_i]);
+  if (this->params_.time_specific_std_dev)
+  {
+    std_dev = &(this->params_.std_dev[(distribution_i * this->getNumTimesteps() + t) * CONTROL_DIM]);
+  }
+  for (int i = 0; i < CONTROL_DIM; i++)
+  {
+    cost += this->params_.control_cost_coeff[i] * u_fb(i) * u_fb(i) / (std_dev[i] * std_dev[i]);
+  }
+  return 0.5f * lambda * (1.0f - alpha) * cost;
+}
+
+GAUSSIAN_TEMPLATE
+__host__ float GAUSSIAN_CLASS::computeLikelihoodRatioCost(const Eigen::Ref<const control_array>& u,
+                                                          const int sample_index, const int t,
                                                           const int distribution_idx, const float lambda,
                                                           const float alpha)
 {
@@ -638,16 +647,21 @@ __host__ float GAUSSIAN_CLASS::computeLikelihoodRatioCost(const Eigen::Ref<const
   const int mean_index = (distribution_i * this->getNumTimesteps() + t) * CONTROL_DIM;
   float* mean = &(this->means_[mean_index]);
   float* std_dev = &(this->params_.std_dev[CONTROL_DIM * distribution_i]);
+  control_array zero_mean = control_array::Zero();
+  if (sample_index >= (1.0f - this->params_.pure_noise_trajectories_percentage) * this->params_.num_rollouts)
+  {
+    mean = zero_mean.data();
+  }
   if (this->params_.time_specific_std_dev)
   {
     std_dev = &(this->params_.std_dev[(distribution_i * this->getNumTimesteps() + t) * CONTROL_DIM]);
   }
   for (int i = 0; i < CONTROL_DIM; i++)
   {
-    cost += this->params_.control_cost_coeff[i] * mean[i] * (mean[i] + 2.0f * u(i)) /
+    cost += this->params_.control_cost_coeff[i] * mean[i] * (mean[i] - 2.0f * u(i)) /
             (std_dev[i] * std_dev[i]);  // Proper way
   }
-  return cost;
+  return 0.5f * lambda * (1.0f - alpha) * cost;
 }
 
 GAUSSIAN_TEMPLATE

include/mppi/sampling_distributions/gaussian/gaussian.cuh

@@ -125,12 +125,15 @@ public:
                                                 const int t, const int distribution_idx, const float lambda = 1.0,
                                                 const float alpha = 0.0);
 
+  __host__ float computeFeedbackCost(const Eigen::Ref<const control_array>& u_fb, const int t,
+                                     const int distribution_idx, const float lambda = 1.0, const float alpha = 0.0);
+
   __host__ __device__ float computeLikelihoodRatioCost(const float* __restrict__ u, float* __restrict__ theta_d,
                                                        const int sample_index, const int t, const int distribution_idx,
                                                        const float lambda = 1.0, const float alpha = 0.0);
 
-  __host__ float computeLikelihoodRatioCost(const Eigen::Ref<const control_array>& u, const int t,
-                                            const int distribution_idx, const float lambda = 1.0,
+  __host__ float computeLikelihoodRatioCost(const Eigen::Ref<const control_array>& u, const int sample_index,
+                                            const int t, const int distribution_idx, const float lambda = 1.0,
                                             const float alpha = 0.0);
 
   __host__ void copyImportanceSamplerToDevice(const float* importance_sampler, const int& distribution_idx,
@@ -146,6 +149,17 @@ public:
   __host__ void setHostOptimalControlSequence(float* optimal_control_trajectory, const int& distribution_idx,
                                               bool synchronize = true);
 
+  __host__ void setNumDistributions(const int num_distributions, bool synchronize = false)
+  {
+    if (num_distributions > SAMPLING_PARAMS_T::MAX_DISTRIBUTIONS)
+    {
+      this->logger_->error("GaussianParams can't handle more than %d distributions but %d were requested\n",
+                           SAMPLING_PARAMS_T::MAX_DISTRIBUTIONS, num_distributions);
+      throw std::out_of_range("Can't set num distributions higher than allowed in params");
+    }
+    PARENT_CLASS::setNumDistributions(num_distributions, synchronize);
+  }
+
   __host__ void updateDistributionParamsFromDevice(const float* trajectory_weights_d, float normalizer,
                                                    const int& distribution_i, bool synchronize = false) override;
 

include/mppi/sampling_distributions/nln/nln.cu

@@ -74,7 +74,7 @@ void NLN_NOISE::setParams(const SAMPLING_PARAMS_T& params, bool synchronize)
   bool adjusted_variance = false;
   for (int i = 0; i < this->CONTROL_DIM * this->getNumDistributions(); i++)
   {
-    if (this->params.std_dev[i] != params.std_dev[i])
+    if (this->params_.std_dev[i] != params.std_dev[i])
     {
       adjusted_variance = true;
       break;