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;