/*
 * Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 *
 * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
 * property and proprietary rights in and to this material, related
 * documentation and any modifications thereto. Any use, reproduction,
 * disclosure or distribution of this material and related documentation
 * without an express license agreement from NVIDIA CORPORATION or
 * its affiliates is strictly prohibited.
 */

#include <c10/cuda/CUDAStream.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <torch/extension.h>

#include "helper_math.h"
#include <cmath>
#include <cuda_fp16.h>
#include <vector>
#define MAX_H 100
#define BWD_DIFF -1
#define CENTRAL_DIFF 0

namespace Curobo
{
  namespace TensorStep
  {
    template<typename scalar_t>
    __global__ void position_clique_loop_kernel(
      scalar_t *out_position_mem, scalar_t *out_velocity_mem,
      scalar_t *out_acceleration_mem, scalar_t *out_jerk_mem,
      const scalar_t *u_position, const scalar_t *start_position,
      const scalar_t *start_velocity, const scalar_t *start_acceleration,
      const scalar_t *traj_dt, const int batch_size, const int horizon,
      const int dof)
    {
      // there are batch * horizon * dof threads:
      const int tid   = blockDim.x * blockIdx.x + threadIdx.x;
      const int b_idx = tid / (dof);
      const int d_idx = (tid - b_idx * dof);

      if ((b_idx >= batch_size) || (d_idx >= dof))
      {
        return;
      }

      const float dt =
        traj_dt[0]; // assume same dt across traj TODO: Implement variable dt

      // read start state:
      float u_arr[MAX_H];
      float out_pos[MAX_H], out_vel[MAX_H], out_acc[MAX_H], out_jerk[MAX_H];
      const int b_addrs = b_idx * horizon * dof;

#pragma unroll

      for (int i = 0; i < horizon; i++)
      {
        u_arr[i]    = u_position[b_addrs + i * dof + d_idx];
        out_pos[i]  = 0;
        out_vel[i]  = 0;
        out_acc[i]  = 0;
        out_jerk[i] = 0;
      }

      out_pos[0]  = start_position[b_idx * dof + d_idx];
      out_vel[0]  = start_velocity[b_idx * dof + d_idx];
      out_acc[0]  = start_acceleration[b_idx * dof + d_idx];
      out_jerk[0] = 0.0;

      for (int h_idx = 1; h_idx < horizon; h_idx++)
      {
        // read actions: batch, horizon
        out_pos[h_idx] = u_arr[h_idx - 1];
        out_vel[h_idx] = (out_pos[h_idx] - out_pos[h_idx - 1]) * dt;  // 1 - 0
        out_acc[h_idx] =
          (out_vel[h_idx] - out_vel[h_idx - 1]) * dt;                 // 2 - 2.0 * 1 + 1
        out_jerk[h_idx] = (out_acc[h_idx] - out_acc[h_idx - 1]) * dt; // -1 3 -3 1
      }

      // write out:
      for (int h_idx = 0; h_idx < horizon; h_idx++)
      {
        out_position_mem[b_idx * horizon * dof + h_idx * dof + d_idx] =
          out_pos[h_idx]; // new_position;
        out_velocity_mem[b_idx * horizon * dof + h_idx * dof + d_idx] =
          out_vel[h_idx];
        out_acceleration_mem[b_idx * horizon * dof + h_idx * dof + d_idx] =
          out_acc[h_idx];
        out_jerk_mem[b_idx * horizon * dof + h_idx * dof + d_idx] = out_jerk[h_idx];
      }
    }

    template<typename scalar_t>
    __device__ __forceinline__ void compute_backward_difference(scalar_t       *out_position_mem,
                                                                scalar_t       *out_velocity_mem,
                                                                scalar_t       *out_acceleration_mem,
                                                                scalar_t       *out_jerk_mem,
                                                                const scalar_t *u_position,
                                                                const scalar_t *start_position,
                                                                const scalar_t *start_velocity,
                                                                const scalar_t *start_acceleration,
                                                                const scalar_t *traj_dt,
                                                                const int       batch_size,
                                                                const int       horizon,
                                                                const int       dof,
                                                                const int       b_idx,
                                                                const int       h_idx,
                                                                const int       d_idx,
                                                                const int       b_offset)
    {
      const float dt =
        traj_dt[0]; // assume same dt across traj TODO: Implement variable dt

      // read start state:
      float in_pos[4];
      float st_pos = 0.0, st_vel = 0.0, st_acc = 0.0;

  #pragma unroll 4

      for (int i = 0; i < 4; i++)
      {
        in_pos[i] = 0.0;
      }

      float out_pos = 0.0, out_vel = 0.0, out_acc = 0.0, out_jerk = 0.0;
      const int b_addrs = b_idx * horizon * dof;

      if (h_idx < 4)
      {
        st_pos = start_position[b_offset * dof + d_idx];
        st_vel = start_velocity[b_offset * dof + d_idx];
        st_acc = start_acceleration[b_offset * dof + d_idx];
      }

      if (h_idx == 0)
      {
        out_pos = st_pos;
        out_vel = st_vel;
        out_acc = st_acc; //
      }
      else
      {
        if (h_idx == 1)
        {
          for (int i = 3; i < 4; i++)
          {
            in_pos[i] = u_position[b_addrs + (h_idx - 1 - 3 + i) * dof + d_idx];
          }


          in_pos[0] = st_pos - 2.0 * (st_vel - 0.5 * st_acc * dt) * dt;

          in_pos[1] = st_pos - (st_vel - 0.5 * st_acc * dt) * dt;

          in_pos[2] = st_pos;
        }
        else if (h_idx == 2)
        {
          for (int i = 2; i < 4; i++)
          {
            in_pos[i] = u_position[b_addrs + (h_idx - 1 - 3 + i) * dof + d_idx];
          }

          in_pos[0] = st_pos - (st_vel - 0.5 * st_acc * dt) * dt;

          in_pos[1] = st_pos;
        }
        else if (h_idx == 3)
        {
          for (int i = 1; i < 4; i++)
          {
            in_pos[i] = u_position[b_addrs + (h_idx - 1 - 3 + i) * dof + d_idx];
          }
          in_pos[0] = st_pos;
        }
        else // h_idx >= 4

        {
          for (int i = 0; i < 4; i++)
          {
            in_pos[i] = u_position[b_addrs + (h_idx - 1 - 3 + i) * dof + d_idx];
          }
        }

        out_pos  = in_pos[3];
        out_vel  = (-in_pos[2] + in_pos[3]) * dt;
        out_acc  = (in_pos[1] - 2 * in_pos[2] + in_pos[3]) * dt * dt;
        out_jerk = (-in_pos[0] + 3 * in_pos[1] - 3 * in_pos[2] + in_pos[3]) * dt * dt * dt;
      }


      // write out:
      out_position_mem[b_idx * horizon * dof + h_idx * dof + d_idx] =
        out_pos; // new_position;
      out_velocity_mem[b_idx * horizon * dof + h_idx * dof + d_idx]     = out_vel;
      out_acceleration_mem[b_idx * horizon * dof + h_idx * dof + d_idx] = out_acc;
      out_jerk_mem[b_idx * horizon * dof + h_idx * dof + d_idx]         = out_jerk;
    }

    template<typename scalar_t>
    __device__ __forceinline__ void compute_central_difference_v0(scalar_t       *out_position_mem,
                                                                  scalar_t       *out_velocity_mem,
                                                                  scalar_t       *out_acceleration_mem,
                                                                  scalar_t       *out_jerk_mem,
                                                                  const scalar_t *u_position,
                                                                  const scalar_t *start_position,
                                                                  const scalar_t *start_velocity,
                                                                  const scalar_t *start_acceleration,
                                                                  const scalar_t *traj_dt,
                                                                  const int       batch_size,
                                                                  const int       horizon,
                                                                  const int       dof,
                                                                  const int       b_idx,
                                                                  const int       h_idx,
                                                                  const int       d_idx,
                                                                  const int       b_offset)
    {
      const float dt = traj_dt[0]; // assume same dt across traj TODO: Implement variable dt
      // dt here is actually 1/dt;

      // read start state:
      float out_pos = 0.0, out_vel = 0.0, out_acc = 0.0, out_jerk = 0.0;
      float st_pos = 0.0, st_vel = 0.0, st_acc = 0.0;

      const int b_addrs = b_idx * horizon * dof;
      float     in_pos[5]; // create a 5 value scalar

  #pragma unroll 5

      for (int i = 0; i < 5; i++)
      {
        in_pos[i] = 0.0;
      }

      if (h_idx < 4)
      {
        st_pos = start_position[b_offset * dof + d_idx];
        st_vel = start_velocity[b_offset * dof + d_idx];
        st_acc = start_acceleration[b_offset * dof + d_idx];
      }

      if (h_idx == 0)
      {
        out_pos = st_pos;
        out_vel = st_vel;
        out_acc = st_acc;
      }
      else if (h_idx < horizon - 2)
      {
        if (h_idx == 1)
        {
          in_pos[0] = st_pos - dt * (st_vel - (0.5 * st_acc * dt)); // start -1, start, u0, u1
          in_pos[1] = st_pos;
          in_pos[2] = u_position[b_addrs + (h_idx - 1) * dof + d_idx];
          in_pos[3] = u_position[b_addrs + (h_idx - 1 + 1) * dof + d_idx];
          in_pos[4] = u_position[b_addrs + (h_idx - 1 + 2) * dof + d_idx];
        }
        else if (h_idx == 2)
        {
          in_pos[0] = start_position[b_offset * dof + d_idx];
          in_pos[1] = u_position[b_addrs + (h_idx - 1 - 1) * dof + d_idx];
          in_pos[2] = u_position[b_addrs + (h_idx - 1) * dof + d_idx];
          in_pos[3] = u_position[b_addrs + (h_idx - 1 + 1) * dof + d_idx];
          in_pos[4] = u_position[b_addrs + (h_idx - 1 + 2) * dof + d_idx]; \

        }


        else if (h_idx > 2)
        {
          in_pos[0] = u_position[b_addrs + (h_idx - 1 - 2) * dof + d_idx];
          in_pos[1] = u_position[b_addrs + (h_idx - 1 - 1) * dof + d_idx];
          in_pos[2] = u_position[b_addrs + (h_idx - 1) * dof + d_idx];
          in_pos[3] = u_position[b_addrs + (h_idx - 1 + 1) * dof + d_idx];
          in_pos[4] = u_position[b_addrs + (h_idx - 1 + 2) * dof + d_idx];
        }
        out_pos  = in_pos[2];
        out_vel  = (0.5 * in_pos[3] - 0.5 * in_pos[1]) * dt;
        out_acc  = (in_pos[3] + in_pos[1] - 2 * in_pos[2]) * dt * dt;
        out_jerk = ((-0.5) * in_pos[0] + in_pos[1]  - in_pos[3] + (0.5) * in_pos[4]) *
                   (dt * dt * dt);
      }
      else if (h_idx == horizon - 2)
      {
        // use backward difference for jerk

        in_pos[0] = u_position[b_addrs + (h_idx - 1 - 3) * dof + d_idx];
        in_pos[1] = u_position[b_addrs + (h_idx - 1 - 2) * dof + d_idx];
        in_pos[2] = u_position[b_addrs + (h_idx - 1 - 1) * dof + d_idx];
        in_pos[3] = u_position[b_addrs + (h_idx - 1) * dof + d_idx];
        in_pos[4] = u_position[b_addrs + (h_idx - 1 + 1) * dof + d_idx];

        out_pos  = in_pos[3];
        out_vel  = (0.5 * in_pos[4] - 0.5 * in_pos[2]) * dt;
        out_acc  = (in_pos[4] + in_pos[2] - 2 * in_pos[3]) * dt * dt;
        out_jerk = (-1 * in_pos[0] + 3 * in_pos[1] - 3 * in_pos[2] + in_pos[3]) * dt * dt * dt;
      }
      else if (h_idx == horizon - 1)
      { // use backward difference for vel, acc, jerk
        for (int i = 0; i < 4; i++)
        {
          in_pos[i] = u_position[b_addrs + (h_idx - 1 - 3 + i) * dof + d_idx];
        }
        out_pos  = in_pos[3];
        out_vel  = (-in_pos[2] + in_pos[3]) * dt;
        out_acc  = (in_pos[1] - 2 * in_pos[2] + in_pos[3]) * dt * dt;
        out_jerk = (-in_pos[0] + 3 * in_pos[1] - 3 * in_pos[2] + in_pos[3]) * dt * dt * dt;
      }

      // write out:
      out_position_mem[b_idx * horizon * dof + h_idx * dof + d_idx]     = out_pos;
      out_velocity_mem[b_idx * horizon * dof + h_idx * dof + d_idx]     = out_vel;
      out_acceleration_mem[b_idx * horizon * dof + h_idx * dof + d_idx] = out_acc;
      out_jerk_mem[b_idx * horizon * dof + h_idx * dof + d_idx]         = out_jerk;
    }

    template<typename scalar_t>
    __device__ __forceinline__ void compute_central_difference(scalar_t       *out_position_mem,
                                                               scalar_t       *out_velocity_mem,
                                                               scalar_t       *out_acceleration_mem,
                                                               scalar_t       *out_jerk_mem,
                                                               const scalar_t *u_position,
                                                               const scalar_t *start_position,
                                                               const scalar_t *start_velocity,
                                                               const scalar_t *start_acceleration,
                                                               const scalar_t *traj_dt,
                                                               const int       batch_size,
                                                               const int       horizon,
                                                               const int       dof,
                                                               const int       b_idx,
                                                               const int       h_idx,
                                                               const int       d_idx,
                                                               const int       b_offset)
    {
      const float dt = traj_dt[0]; // assume same dt across traj TODO: Implement variable dt
      // dt here is actually 1/dt;
      const float dt_inv  = 1.0 / dt;
      const float st_jerk = 0.0;   // Note: start jerk can also be passed from global memory
      // read start state:
      float out_pos = 0.0, out_vel = 0.0, out_acc = 0.0, out_jerk = 0.0;
      float st_pos = 0.0, st_vel = 0.0, st_acc = 0.0;

      const int   b_addrs_action = b_idx * (horizon - 4) * dof;
      float       in_pos[5]; // create a 5 value scalar

  #pragma unroll 5

      for (int i = 0; i < 5; i++)
      {
        in_pos[i] = 0.0;
      }

      if (h_idx < 5)
      {
        st_pos = start_position[b_offset * dof + d_idx];
        st_vel = start_velocity[b_offset * dof + d_idx];
        st_acc = start_acceleration[b_offset * dof + d_idx];
      }

      if ((h_idx > 3) && (h_idx < horizon - 4))
      {
        in_pos[0] = u_position[b_addrs_action + (h_idx - 4) * dof + d_idx];
        in_pos[1] = u_position[b_addrs_action + (h_idx - 3) * dof + d_idx];
        in_pos[2] = u_position[b_addrs_action + (h_idx - 2) * dof + d_idx];
        in_pos[3] = u_position[b_addrs_action + (h_idx - 1) * dof + d_idx];
        in_pos[4] = u_position[b_addrs_action + (h_idx) * dof + d_idx];
      }


      else if (h_idx == 0)
      {
        in_pos[0] = (3.0f / 2) *
                    (-1 * st_acc * (dt_inv * dt_inv) - (dt_inv * dt_inv * dt_inv) * st_jerk) -
                    3.0f * dt_inv *
                    st_vel + st_pos;
        in_pos[1] = -2.0f * st_acc * dt_inv * dt_inv - (4.0f / 3) * dt_inv * dt_inv * dt_inv *
                    st_jerk - 2.0 * dt_inv * st_vel + st_pos;
        in_pos[2] = -(3.0f / 2) * st_acc * dt_inv * dt_inv - (7.0f / 6) * dt_inv * dt_inv * dt_inv *
                    st_jerk - dt_inv * st_vel + st_pos;
        in_pos[3] = st_pos;
        in_pos[4] = u_position[b_addrs_action + (h_idx) * dof + d_idx];
      }

      else if (h_idx == 1)
      {
        in_pos[0] = -2.0f * st_acc * dt_inv * dt_inv - (4.0f / 3) * dt_inv * dt_inv * dt_inv *
                    st_jerk - 2.0 * dt_inv * st_vel + st_pos;
        in_pos[1] = -(3.0f / 2) * st_acc * dt_inv * dt_inv - (7.0f / 6) * dt_inv * dt_inv * dt_inv *
                    st_jerk - dt_inv * st_vel + st_pos;


        in_pos[2] = st_pos;
        in_pos[3] = u_position[b_addrs_action + (h_idx - 1) * dof + d_idx];
        in_pos[4] = u_position[b_addrs_action + (h_idx) * dof + d_idx];
      }

      else if (h_idx == 2)
      {
        in_pos[0] = -(3.0f / 2) * st_acc * dt_inv * dt_inv - (7.0f / 6) * dt_inv * dt_inv * dt_inv *
                    st_jerk - dt_inv * st_vel + st_pos;
        in_pos[1] = st_pos;
        in_pos[2] = u_position[b_addrs_action + (h_idx - 2) * dof + d_idx];
        in_pos[3] = u_position[b_addrs_action + (h_idx - 1) * dof + d_idx];
        in_pos[4] = u_position[b_addrs_action + (h_idx) * dof + d_idx];
      }
      else if (h_idx == 3)
      {
        in_pos[0] = st_pos;
        in_pos[1] = u_position[b_addrs_action + (h_idx - 3) * dof + d_idx];
        in_pos[2] = u_position[b_addrs_action + (h_idx - 2) * dof + d_idx];
        in_pos[3] = u_position[b_addrs_action + (h_idx - 1) * dof + d_idx];
        in_pos[4] = u_position[b_addrs_action + (h_idx) * dof + d_idx];
      }

      else if (h_idx == horizon - 4)
      {
        in_pos[0] = u_position[b_addrs_action + (h_idx - 4) * dof + d_idx];
        in_pos[1] = u_position[b_addrs_action + (h_idx - 3) * dof + d_idx];
        in_pos[2] = u_position[b_addrs_action + (h_idx - 2) * dof + d_idx];
        in_pos[3] = u_position[b_addrs_action + (h_idx - 1) * dof + d_idx];
        in_pos[4] = in_pos[3]; // in_pos[3]; //u_position[b_addrs_action + (h_idx - 1 + 2) * dof +
                               // d_idx];
      }

      else if (h_idx == horizon - 3)
      {
        in_pos[0] = u_position[b_addrs_action + (h_idx - 4) * dof + d_idx];
        in_pos[1] = u_position[b_addrs_action + (h_idx - 3) * dof + d_idx];
        in_pos[2] = u_position[b_addrs_action + (h_idx - 2) * dof + d_idx];
        in_pos[3] = in_pos[2]; // u_position[b_addrs_action + (h_idx - 1 + 1) * dof + d_idx];
        in_pos[4] = in_pos[2]; // in_pos[3]; //u_position[b_addrs_action + (h_idx - 1 + 2) * dof +
                               // d_idx];
      }
      else if (h_idx == horizon - 2)
      {
        in_pos[0] = u_position[b_addrs_action + (h_idx - 4) * dof + d_idx];
        in_pos[1] = u_position[b_addrs_action + (h_idx - 3) * dof + d_idx];
        in_pos[2] = in_pos[1];
        in_pos[3] = in_pos[1]; // u_position[b_addrs_action + (h_idx - 1 + 1) * dof + d_idx];
        in_pos[4] = in_pos[1]; // in_pos[3]; //u_position[b_addrs_action + (h_idx - 1 + 2) * dof +
                               // d_idx];
      }

      else if (h_idx == horizon - 1)
      {
        in_pos[0] = u_position[b_addrs_action + (h_idx - 4) * dof + d_idx];
        in_pos[1] = in_pos[0];
        in_pos[2] = in_pos[0]; // u_position[b_addrs_action + (h_idx - 1 ) * dof + d_idx];
        in_pos[3] = in_pos[0]; // u_position[b_addrs_action + (h_idx - 1 + 1) * dof + d_idx];
        in_pos[4] = in_pos[0]; // in_pos[3]; //u_position[b_addrs_action + (h_idx - 1 + 2) * dof +
                               // d_idx];
      }
      out_pos = in_pos[2];

      // out_vel = (0.5 * in_pos[3] - 0.5 * in_pos[1]) * dt;
      out_vel =
        ((0.083333333f) * in_pos[0] - (0.666666667f) * in_pos[1] + (0.666666667f) * in_pos[3] +
         (-0.083333333f) * in_pos[4]) * dt;

      // out_acc = (in_pos[3] + in_pos[1] - 2.0 * in_pos[2]) * dt * dt;
      out_acc =
        ((-0.083333333f) * in_pos[0] + (1.333333333f) * in_pos[1] + (-2.5f) * in_pos[2] +
         (1.333333333f) * in_pos[3] + (-0.083333333f) * in_pos[4]) * dt * dt;
      out_jerk = ((-0.5f) * in_pos[0] + in_pos[1]  - in_pos[3] + (0.5f) * in_pos[4]) *
                 (dt * dt * dt);

      // write out:
      out_position_mem[b_idx * horizon * dof + h_idx * dof + d_idx]     = out_pos;
      out_velocity_mem[b_idx * horizon * dof + h_idx * dof + d_idx]     = out_vel;
      out_acceleration_mem[b_idx * horizon * dof + h_idx * dof + d_idx] = out_acc;
      out_jerk_mem[b_idx * horizon * dof + h_idx * dof + d_idx]         = out_jerk;
    }

    template<typename scalar_t, int8_t mode>
    __global__ void position_clique_loop_kernel2(
      scalar_t *out_position_mem, scalar_t *out_velocity_mem,
      scalar_t *out_acceleration_mem, scalar_t *out_jerk_mem,
      const scalar_t *u_position, const scalar_t *start_position,
      const scalar_t *start_velocity, const scalar_t *start_acceleration,
      const scalar_t *traj_dt, const int batch_size, const int horizon,
      const int dof)
    {
      const int tid = blockDim.x * blockIdx.x + threadIdx.x;

      // number of threads = batch_size * dof * horizon;
      const int h_idx = tid % horizon;
      const int d_idx = (tid / horizon) % dof;
      const int b_idx = tid / (dof * horizon);

      if (tid >= batch_size * dof * horizon)
      {
        return;
      }


      const int b_offset = b_idx;

      if (mode == BWD_DIFF)
      {
        compute_backward_difference(out_position_mem,
                                    out_velocity_mem,
                                    out_acceleration_mem,
                                    out_jerk_mem,
                                    u_position,
                                    start_position,
                                    start_velocity,
                                    start_acceleration,
                                    traj_dt,
                                    batch_size,
                                    horizon,
                                    dof,
                                    b_idx,
                                    h_idx,
                                    d_idx,
                                    b_offset);
      }
      else if (mode == CENTRAL_DIFF)
      {
        compute_central_difference(out_position_mem,
                                   out_velocity_mem,
                                   out_acceleration_mem,
                                   out_jerk_mem,
                                   u_position,
                                   start_position,
                                   start_velocity,
                                   start_acceleration,
                                   traj_dt,
                                   batch_size,
                                   horizon,
                                   dof,
                                   b_idx,
                                   h_idx,
                                   d_idx,
                                   b_offset);
      }
      else
      {
        assert(false);
      }
    }

    template<typename scalar_t, int8_t mode>
    __global__ void position_clique_loop_idx_kernel2(
      scalar_t *out_position_mem, scalar_t *out_velocity_mem,
      scalar_t *out_acceleration_mem, scalar_t *out_jerk_mem,
      const scalar_t *u_position, const scalar_t *start_position,
      const scalar_t *start_velocity, const scalar_t *start_acceleration,
      const int32_t *start_idx, const scalar_t *traj_dt, const int batch_size,
      const int horizon, const int dof)
    {
      const int tid = blockDim.x * blockIdx.x + threadIdx.x;

      // number of threads = batch_size * dof * horizon;
      const int h_idx = tid % horizon;
      const int d_idx = (tid / horizon) % dof;
      const int b_idx = tid / (dof * horizon);

      if (tid >= batch_size * dof * horizon)
      {
        return;
      }


      const int b_offset = start_idx[b_idx];

      if (mode == BWD_DIFF)
      {
        compute_backward_difference(out_position_mem,
                                    out_velocity_mem,
                                    out_acceleration_mem,
                                    out_jerk_mem,
                                    u_position,
                                    start_position,
                                    start_velocity,
                                    start_acceleration,
                                    traj_dt,
                                    batch_size,
                                    horizon,
                                    dof,
                                    b_idx,
                                    h_idx,
                                    d_idx,
                                    b_offset);
      }
      else if (mode == CENTRAL_DIFF)
      {
        compute_central_difference(out_position_mem,
                                   out_velocity_mem,
                                   out_acceleration_mem,
                                   out_jerk_mem,
                                   u_position,
                                   start_position,
                                   start_velocity,
                                   start_acceleration,
                                   traj_dt,
                                   batch_size,
                                   horizon,
                                   dof,
                                   b_idx,
                                   h_idx,
                                   d_idx,
                                   b_offset);
      }
      else
      {
        assert(false);
      }
    }

    template<typename scalar_t>
    __global__ void backward_position_clique_loop_kernel(
      scalar_t *out_grad_position, const scalar_t *grad_position,
      const scalar_t *grad_velocity, const scalar_t *grad_acceleration,
      const scalar_t *grad_jerk, const scalar_t *traj_dt, const int batch_size,
      const int horizon, const int dof)
    {
      const int tid   = blockDim.x * blockIdx.x + threadIdx.x;
      const int b_idx = tid / (dof);
      const int d_idx = (tid - b_idx * dof);

      if ((b_idx >= batch_size) || (d_idx >= dof))
      {
        return;
      }
      const int b_addrs = b_idx * horizon * dof;

      // read gradients:
      float g_pos[MAX_H];
      float g_vel[MAX_H];
      float g_acc[MAX_H];
      float g_jerk[MAX_H];
      const float dt   = traj_dt[0];
      const float dt_2 = dt * dt;      // dt * dt;
      const float dt_3 = dt * dt * dt; // dt * dt * dt;

      // not used index == 0
      g_pos[0]  = 0.0;
      g_vel[0]  = 0.0;
      g_acc[0]  = 0.0;
      g_jerk[0] = 0.0;
#pragma unroll

      for (int h_idx = 1; h_idx < horizon; h_idx++)
      {
        g_pos[h_idx]  = grad_position[b_addrs + (h_idx) * dof + d_idx];
        g_vel[h_idx]  = grad_velocity[b_addrs + (h_idx) * dof + d_idx];
        g_acc[h_idx]  = grad_acceleration[b_addrs + (h_idx) * dof + d_idx];
        g_jerk[h_idx] = grad_jerk[b_addrs + (h_idx) * dof + d_idx];
      }
#pragma unroll

      for (int i = 0; i < 4; i++)
      {
        g_vel[horizon + i]  = 0.0;
        g_acc[horizon + i]  = 0.0;
        g_jerk[horizon + i] = 0.0;
      }

      // compute gradient and sum
      for (int h_idx = 0; h_idx < horizon - 1; h_idx++)
      {
        g_pos[h_idx + 1] +=
          ((g_vel[h_idx + 1] - g_vel[h_idx + 2]) * dt +
           (g_acc[h_idx + 1] - 2 * g_acc[h_idx + 2] + g_acc[h_idx + 3]) * dt_2 +
           (g_jerk[h_idx + 1] - 3 * g_jerk[h_idx + 2] + 3 * g_jerk[h_idx + 3] -
            g_jerk[h_idx + 4]) *
           dt_3);
      }

      // write out:
      for (int h_idx = 0; h_idx < horizon - 1; h_idx++)
      {
        out_grad_position[b_addrs + (h_idx) * dof + d_idx] = g_pos[h_idx + 1];
      }
      out_grad_position[b_addrs + (horizon - 1) * dof + d_idx] = 0.0;
    }

    template<typename scalar_t>
    __global__ void backward_position_clique_loop_backward_difference_kernel2(
      scalar_t *out_grad_position, const scalar_t *grad_position,
      const scalar_t *grad_velocity, const scalar_t *grad_acceleration,
      const scalar_t *grad_jerk, const scalar_t *traj_dt, const int batch_size,
      const int horizon, const int dof)
    {
      const int tid = blockDim.x * blockIdx.x + threadIdx.x;

      // number of threads = batch_size * dof * horizon;
      const int h_idx = tid % horizon;
      const int d_idx = (tid / horizon) % dof;
      const int b_idx = tid / (dof * horizon);

      if (tid >= batch_size * dof * horizon)
      {
        return;
      }
      const int b_addrs = b_idx * horizon * dof;

      if (h_idx == 0)
      {
        return;
      }


      const float dt = traj_dt[0];

      // read gradients:
      float g_pos    = 0.0;
      float out_grad = 0.0;
      float g_vel[4];
      float g_acc[4];
      float g_jerk[4];

  #pragma unroll 4

      for (int i = 0; i < 4; i++)
      {
        g_vel[i]  = 0.0;
        g_acc[i]  = 0.0;
        g_jerk[i] = 0.0;
      }

      int hid = h_idx; // + 1;

      g_pos = grad_position[b_addrs + (hid) * dof + d_idx];


      if (hid < horizon - 3)
      {
        for (int i = 0; i < 4; i++)
        {
          g_vel[i]  = grad_velocity[b_addrs + (hid + i) * dof + d_idx];
          g_acc[i]  = grad_acceleration[b_addrs + (hid + i) * dof + d_idx];
          g_jerk[i] = grad_jerk[b_addrs + (hid + i) * dof + d_idx];
        }
      }
      else if (hid == horizon - 3)
      {
        for (int i = 0; i < 3; i++)
        {
          g_vel[i]  = grad_velocity[b_addrs + (hid + i) * dof + d_idx];
          g_acc[i]  = grad_acceleration[b_addrs + (hid + i) * dof + d_idx];
          g_jerk[i] = grad_jerk[b_addrs + (hid + i) * dof + d_idx];
        }
      }


      else if (hid == horizon - 2)
      {
        for (int i = 0; i < 2; i++)
        {
          g_vel[i]  = grad_velocity[b_addrs + (hid + i) * dof + d_idx];
          g_acc[i]  = grad_acceleration[b_addrs + (hid + i) * dof + d_idx];
          g_jerk[i] = grad_jerk[b_addrs + (hid + i) * dof + d_idx];
        }
      }
      else if (hid == horizon - 1)
      {
        for (int i = 0; i < 1; i++)
        {
          g_vel[i]  = grad_velocity[b_addrs + (hid + i) * dof + d_idx];
          g_acc[i]  = grad_acceleration[b_addrs + (hid + i) * dof + d_idx];
          g_jerk[i] = grad_jerk[b_addrs + (hid + i) * dof + d_idx];
        }
      }


      out_grad = (g_pos +
                  (g_vel[0] -  g_vel[1]) * dt +
                  (g_acc[0] - 2 * g_acc[1] + g_acc[2]) * dt * dt +
                  (g_jerk[0] - 3 * g_jerk[1] + 3 * g_jerk[2] - g_jerk[3]) * dt * dt * dt);


      // write out:
      out_grad_position[b_addrs + (h_idx - 1) * dof + d_idx] = out_grad;
    }

    template<typename scalar_t>
    __global__ void backward_position_clique_loop_central_difference_kernel2(
      scalar_t *out_grad_position, const scalar_t *grad_position,
      const scalar_t *grad_velocity, const scalar_t *grad_acceleration,
      const scalar_t *grad_jerk, const scalar_t *traj_dt, const int batch_size,
      const int horizon, const int dof)
    {
      const int tid = blockDim.x * blockIdx.x + threadIdx.x;

      // number of threads = batch_size * dof * horizon;
      const int h_idx = tid % horizon;
      const int d_idx = (tid / horizon) % dof;
      const int b_idx = tid / (dof * horizon);

      if (tid >= batch_size * dof * horizon)
      {
        return;
      }
      const int b_addrs        = b_idx * horizon * dof;
      const int b_addrs_action = b_idx * (horizon - 4) * dof;

      if ((h_idx < 2) || (h_idx >= horizon - 2))
      {
        return;
      }

      const float dt = traj_dt[0];

      // read gradients:
      // float g_pos= 0.0;
      float out_grad = 0.0;
      float g_pos[3];
      float g_vel[5];
      float g_acc[5];
      float g_jerk[5];

  #pragma unroll 5

      for (int i = 0; i < 5; i++)
      {
        g_vel[i]  = 0.0;
        g_acc[i]  = 0.0;
        g_jerk[i] = 0.0;
      }

      const int hid = h_idx;

      g_pos[0] = grad_position[b_addrs + (hid) * dof + d_idx];
      g_pos[1] = 0.0;
      g_pos[2] = 0.0;


      if ((hid > 1) && (h_idx < horizon - 3))
      {
    #pragma unroll

        for (int i = 0; i < 5; i++)
        {
          g_vel[i]  = grad_velocity[b_addrs + ((hid - 2) + i) * dof + d_idx];
          g_acc[i]  = grad_acceleration[b_addrs + ((hid - 2) + i) * dof + d_idx];
          g_jerk[i] = grad_jerk[b_addrs + ((hid - 2) + i) * dof + d_idx];
        }
        out_grad = (g_pos[0] +

// ((-0.5) * g_vel[3] + (0.5) * g_vel[1]) * dt +

                    ((-0.083333333f) * g_vel[0] + (0.666666667f) * g_vel[1] + (-0.666666667f) *
                     g_vel[3] + (0.083333333f) * g_vel[4]) * dt +


                    ((-0.083333333f) * g_acc[0] + (1.333333333f) * g_acc[1] + (-2.5f) * g_acc[2] +
                     (1.333333333f) * g_acc[3] + (-0.083333333f) * g_acc[4]) * dt * dt +

// ( g_acc[3] + g_acc[1] - (2.0) * g_acc[2]) * dt * dt +
                    (-0.5f * g_jerk[0] + g_jerk[1] - g_jerk[3] + 0.5f * g_jerk[4]) * dt * dt * dt);
      }
      else if (hid == horizon - 3)
      {
        // The below can cause oscilatory gradient steps.

        /*
         #pragma unroll
           for (int i=0; i< 5; i++)
           {
           g_vel[i] = grad_velocity[b_addrs + ((hid - 2) + i)*dof + d_idx];
           g_acc[i] = grad_acceleration[b_addrs + ((hid -2) + i)*dof + d_idx];
           g_jerk[i] = grad_jerk[b_addrs + ((hid -2) + i)*dof + d_idx];
           }
         */
        g_pos[1] = grad_position[b_addrs + (hid + 1) * dof + d_idx];
        g_pos[2] = grad_position[b_addrs + (hid + 2) * dof + d_idx];

        out_grad = (g_pos[0] + g_pos[1] + g_pos[2]);

        /* +
           //((0.5) * g_vel[1] + (0.5) * g_vel[2]) * dt +
           ((-0.083333333f) * g_vel[0] + (0.583333333f) * g_vel[1] + (0.583333333f) * g_vel[2] +
              (-0.083333333f) * g_vel[3]) * dt  +
           ((-0.083333333f) * g_acc[0] + (1.25f) * g_acc[1] + (-1.25f) * g_acc[2] + (0.083333333f) *
              g_acc[3]) * dt * dt +
           //( g_acc[1] - g_acc[2]) * dt * dt +
           (0.5f * g_jerk[0] - 0.5f * g_jerk[1] -0.5f * g_jerk[2] + 0.5f * g_jerk[3]) * dt * dt *
              dt);
         */
      }


      // write out:
      out_grad_position[b_addrs_action + (h_idx - 2) * dof + d_idx] = out_grad;
    }

    // for MPPI:

    template<typename scalar_t>
    __global__ void
    acceleration_loop_kernel(scalar_t *out_position_mem, scalar_t *out_velocity_mem,
                             scalar_t *out_acceleration_mem, scalar_t *out_jerk_mem,
                             const scalar_t *u_acc, const scalar_t *start_position,
                             const scalar_t *start_velocity,
                             const scalar_t *start_acceleration,
                             const scalar_t *traj_dt, const int batch_size,
                             const int horizon, const int dof)
    {
      // there are batch * horizon * dof threads:
      const int tid   = blockDim.x * blockIdx.x + threadIdx.x;
      const int b_idx = tid / (dof);
      const int d_idx = (tid - b_idx * dof);

      if ((b_idx >= batch_size) || (d_idx >= dof))
      {
        return;
      }

      // read start state:
      float u_arr[MAX_H], dt[MAX_H];
      float out_pos[MAX_H], out_vel[MAX_H], out_acc[MAX_H], out_jerk[MAX_H];
      const int b_addrs = b_idx * horizon * dof;

#pragma unroll

      for (int i = 0; i < horizon; i++)
      {
        u_arr[i]    = u_acc[b_addrs + i * dof + d_idx];
        dt[i]       = traj_dt[i];
        out_pos[i]  = 0;
        out_vel[i]  = 0;
        out_acc[i]  = 0;
        out_jerk[i] = 0;
      }

      out_pos[0]  = start_position[b_idx * dof + d_idx];
      out_vel[0]  = start_velocity[b_idx * dof + d_idx];
      out_acc[0]  = start_acceleration[b_idx * dof + d_idx];
      out_jerk[0] = 0.0;

      for (int h_idx = 1; h_idx < horizon; h_idx++)
      {
        // do semi implicit euler integration:
        out_acc[h_idx]  = u_arr[h_idx - 1];
        out_vel[h_idx]  = out_vel[h_idx - 1] + out_acc[h_idx] * dt[h_idx];
        out_pos[h_idx]  = out_pos[h_idx - 1] + out_vel[h_idx] * dt[h_idx];
        out_jerk[h_idx] = (out_acc[h_idx] - out_acc[h_idx - 1]) / dt[h_idx];
      }

      // write out:
      for (int h_idx = 0; h_idx < horizon; h_idx++)
      {
        out_position_mem[b_idx * horizon * dof + h_idx * dof + d_idx] =
          out_pos[h_idx]; // new_position;
        out_velocity_mem[b_idx * horizon * dof + h_idx * dof + d_idx] =
          out_vel[h_idx];
        out_acceleration_mem[b_idx * horizon * dof + h_idx * dof + d_idx] =
          out_acc[h_idx];
        out_jerk_mem[b_idx * horizon * dof + h_idx * dof + d_idx] = out_jerk[h_idx];
      }
    }

    template<typename scalar_t>
    __global__ void acceleration_loop_rk2_kernel(
      scalar_t *out_position_mem, scalar_t *out_velocity_mem,
      scalar_t *out_acceleration_mem, scalar_t *out_jerk_mem,
      const scalar_t *u_acc, const scalar_t *start_position,
      const scalar_t *start_velocity, const scalar_t *start_acceleration,
      const scalar_t *traj_dt, const int batch_size, const int horizon,
      const int dof)
    {
      // there are batch * horizon * dof threads:
      const int tid   = blockDim.x * blockIdx.x + threadIdx.x;
      const int b_idx = tid / (dof);
      const int d_idx = (tid - b_idx * dof);

      if ((b_idx >= batch_size) || (d_idx >= dof))
      {
        return;
      }

      // read start state:
      float u_arr[MAX_H], dt[MAX_H];
      float out_pos[MAX_H], out_vel[MAX_H], out_acc[MAX_H], out_jerk[MAX_H];
      const int b_addrs = b_idx * horizon * dof;

#pragma unroll

      for (int i = 0; i < horizon; i++)
      {
        u_arr[i]    = u_acc[b_addrs + i * dof + d_idx];
        dt[i]       = traj_dt[i];
        out_pos[i]  = 0;
        out_vel[i]  = 0;
        out_acc[i]  = 0;
        out_jerk[i] = 0;
      }

      out_pos[0]  = start_position[b_idx * dof + d_idx];
      out_vel[0]  = start_velocity[b_idx * dof + d_idx];
      out_acc[0]  = start_acceleration[b_idx * dof + d_idx];
      out_jerk[0] = 0.0;

      for (int h_idx = 1; h_idx < horizon; h_idx++)
      {
        // do rk2 integration:

        out_acc[h_idx]  = u_arr[h_idx - 1];
        out_jerk[h_idx] = (out_acc[h_idx] - out_acc[h_idx - 1]) / dt[h_idx];
        out_pos[h_idx]  = out_pos[h_idx - 1] + out_vel[h_idx - 1] * dt[h_idx] +
                          0.5 * dt[h_idx] * dt[h_idx] * out_acc[h_idx];
        out_vel[h_idx] = out_vel[h_idx - 1] + 0.5 * dt[h_idx] * out_acc[h_idx];
      }

      // write out:
      for (int h_idx = 0; h_idx < horizon; h_idx++)
      {
        out_position_mem[b_idx * horizon * dof + h_idx * dof + d_idx] =
          out_pos[h_idx]; // new_position;
        out_velocity_mem[b_idx * horizon * dof + h_idx * dof + d_idx] =
          out_vel[h_idx];
        out_acceleration_mem[b_idx * horizon * dof + h_idx * dof + d_idx] =
          out_acc[h_idx];
        out_jerk_mem[b_idx * horizon * dof + h_idx * dof + d_idx] = out_jerk[h_idx];
      }
    }

    template<typename scalar_t>
    __global__ void acceleration_loop_idx_kernel(
      scalar_t *out_position_mem, scalar_t *out_velocity_mem,
      scalar_t *out_acceleration_mem, scalar_t *out_jerk_mem,
      const scalar_t *u_acc, const scalar_t *start_position,
      const scalar_t *start_velocity, const scalar_t *start_acceleration,
      const int32_t *start_idx, const scalar_t *traj_dt, const int batch_size,
      const int horizon, const int dof)
    {
      // there are batch * horizon * dof threads:
      const int tid   = blockDim.x * blockIdx.x + threadIdx.x;
      const int b_idx = tid / (dof);
      const int d_idx = (tid - b_idx * dof);

      if ((b_idx >= batch_size) || (d_idx >= dof))
      {
        return;
      }

      // read start state:
      float u_arr[MAX_H], dt[MAX_H];
      float out_pos[MAX_H], out_vel[MAX_H], out_acc[MAX_H], out_jerk[MAX_H];
      const int b_addrs  = b_idx * horizon * dof;
      const int b_offset = start_idx[b_idx];

#pragma unroll

      for (int i = 0; i < horizon; i++)
      {
        u_arr[i]    = u_acc[b_addrs + i * dof + d_idx];
        dt[i]       = traj_dt[i];
        out_pos[i]  = 0;
        out_vel[i]  = 0;
        out_acc[i]  = 0;
        out_jerk[i] = 0;
      }

      out_pos[0]  = start_position[b_offset * dof + d_idx];
      out_vel[0]  = start_velocity[b_offset * dof + d_idx];
      out_acc[0]  = start_acceleration[b_offset * dof + d_idx];
      out_jerk[0] = 0.0;

      for (int h_idx = 1; h_idx < horizon; h_idx++)
      {
        // do semi implicit euler integration:
        out_acc[h_idx]  = u_arr[h_idx - 1];
        out_vel[h_idx]  = out_vel[h_idx - 1] + out_acc[h_idx] * dt[h_idx];
        out_pos[h_idx]  = out_pos[h_idx - 1] + out_vel[h_idx] * dt[h_idx];
        out_jerk[h_idx] = (out_acc[h_idx] - out_acc[h_idx - 1]) / dt[h_idx];
      }

      // write out:
      for (int h_idx = 0; h_idx < horizon; h_idx++)
      {
        out_position_mem[b_idx * horizon * dof + h_idx * dof + d_idx] =
          out_pos[h_idx]; // new_position;
        out_velocity_mem[b_idx * horizon * dof + h_idx * dof + d_idx] =
          out_vel[h_idx];
        out_acceleration_mem[b_idx * horizon * dof + h_idx * dof + d_idx] =
          out_acc[h_idx];
        out_jerk_mem[b_idx * horizon * dof + h_idx * dof + d_idx] = out_jerk[h_idx];
      }
    }

    template<typename scalar_t>
    __global__ void acceleration_loop_idx_rk2_kernel(
      scalar_t *out_position_mem, scalar_t *out_velocity_mem,
      scalar_t *out_acceleration_mem, scalar_t *out_jerk_mem,
      const scalar_t *u_acc, const scalar_t *start_position,
      const scalar_t *start_velocity, const scalar_t *start_acceleration,
      const int32_t *start_idx, const scalar_t *traj_dt, const int batch_size,
      const int horizon, const int dof)
    {
      // there are batch * horizon * dof threads:
      const int tid   = blockDim.x * blockIdx.x + threadIdx.x;
      const int b_idx = tid / (dof);
      const int d_idx = (tid - b_idx * dof);

      if ((b_idx >= batch_size) || (d_idx >= dof))
      {
        return;
      }

      // read start state:
      float u_arr[MAX_H], dt[MAX_H];
      float out_pos[MAX_H], out_vel[MAX_H], out_acc[MAX_H], out_jerk[MAX_H];
      const int b_addrs  = b_idx * horizon * dof;
      const int b_offset = start_idx[b_idx];

#pragma unroll

      for (int i = 0; i < horizon; i++)
      {
        u_arr[i]    = u_acc[b_addrs + i * dof + d_idx];
        dt[i]       = traj_dt[i];
        out_pos[i]  = 0;
        out_vel[i]  = 0;
        out_acc[i]  = 0;
        out_jerk[i] = 0;
      }

      out_pos[0]  = start_position[b_offset * dof + d_idx];
      out_vel[0]  = start_velocity[b_offset * dof + d_idx];
      out_acc[0]  = start_acceleration[b_offset * dof + d_idx];
      out_jerk[0] = 0.0;

      for (int h_idx = 1; h_idx < horizon; h_idx++)
      {
        // do semi implicit euler integration:

        out_acc[h_idx]  = u_arr[h_idx - 1];
        out_vel[h_idx]  = out_vel[h_idx - 1] + out_acc[h_idx] * dt[h_idx];
        out_pos[h_idx]  = out_pos[h_idx - 1] + out_vel[h_idx] * dt[h_idx];
        out_jerk[h_idx] = (out_acc[h_idx] - out_acc[h_idx - 1]) / dt[h_idx];
      }

      // write out:
      for (int h_idx = 0; h_idx < horizon; h_idx++)
      {
        out_position_mem[b_idx * horizon * dof + h_idx * dof + d_idx] =
          out_pos[h_idx]; // new_position;
        out_velocity_mem[b_idx * horizon * dof + h_idx * dof + d_idx] =
          out_vel[h_idx];
        out_acceleration_mem[b_idx * horizon * dof + h_idx * dof + d_idx] =
          out_acc[h_idx];
        out_jerk_mem[b_idx * horizon * dof + h_idx * dof + d_idx] = out_jerk[h_idx];
      }
    }

    // Not used

    template<typename scalar_t>
    __global__ void position_clique_kernel(
      scalar_t *out_position, scalar_t *out_velocity, scalar_t *out_acceleration,
      scalar_t *out_jerk, const scalar_t *u_position,
      const scalar_t *start_position, const scalar_t *start_velocity,
      const scalar_t *start_acceleration, const scalar_t *traj_dt,
      const int batch_size, const int horizon, const int dof)
    {
      // there are batch * horizon * dof threads:
      const int tid   = blockDim.x * blockIdx.x + threadIdx.x;
      const int b_idx = tid / (horizon * dof);
      const int h_idx = (tid - b_idx * (horizon * dof)) / dof;
      const int d_idx = (tid - b_idx * horizon * dof - h_idx * dof);

      if ((b_idx >= batch_size) || (h_idx >= horizon) || (d_idx >= dof))
      {
        return;
      }

      float new_position, new_velocity, new_acceleration, new_jerk;
      const float dt = traj_dt[h_idx];

      // read actions: batch, horizon
      if (h_idx == 0)
      {
        new_position     = start_position[b_idx * dof + d_idx];
        new_velocity     = start_velocity[b_idx * dof + d_idx];
        new_acceleration = start_acceleration[b_idx * dof + d_idx];
        new_jerk         = 0.0;
      }
      else if (h_idx == 1)
      {
        float2 u_clique = make_float2(
          start_position[b_idx * dof + d_idx],
          u_position[b_idx * horizon * dof + (h_idx - 1) * dof + d_idx]);
        new_position     = u_clique.y;
        new_velocity     = (u_clique.y - u_clique.x) * dt;    // 1 - 0
        new_acceleration =
          (u_clique.y - u_clique.x) * dt * dt;                // 2 - 2.0 * 1 + 1
        new_jerk = (u_clique.y - u_clique.x) *  dt * dt * dt; // -1 3 -3 1
      }
      else if (h_idx == 2)
      {
        float3 u_clique = make_float3(
          start_position[b_idx * dof + d_idx],
          u_position[b_idx * horizon * dof + (h_idx - 2) * dof + d_idx],
          u_position[b_idx * horizon * dof + (h_idx - 1) * dof + d_idx]);

        new_position     = u_clique.z;
        new_velocity     = (u_clique.z - u_clique.y) * dt; // 1 - 0
        new_acceleration = (u_clique.x - 2 * u_clique.y + u_clique.z) *
                           dt * dt;                        // 2 - 2.0 * 1 + 1
        new_jerk = (2 * u_clique.x - 3 * u_clique.y + u_clique.z) *
                   dt * dt * dt;                           // -1 3 -3 1
      }
      else if (h_idx == 3)
      {
        float4 u_clique = make_float4(
          start_position[b_idx * dof + d_idx],
          u_position[b_idx * horizon * dof + (h_idx - 3) * dof + d_idx],
          u_position[b_idx * horizon * dof + (h_idx - 2) * dof + d_idx],
          u_position[b_idx * horizon * dof + (h_idx - 1) * dof + d_idx]);
        new_position     = u_clique.w;
        new_velocity     = (u_clique.w - u_clique.z) * dt; // 1 - 0
        new_acceleration = (u_clique.y - 2 * u_clique.z + u_clique.w) *
                           dt * dt;                        // 2 - 2.0 * 1 + 1
        new_jerk =
          (-1.0 * u_clique.x + 3 * u_clique.y - 3 * u_clique.z + u_clique.w) *
          dt * dt * dt;
      }
      else
      {
        float4 u_clique = make_float4(
          u_position[b_idx * horizon * dof + (h_idx - 4) * dof + d_idx],
          u_position[b_idx * horizon * dof + (h_idx - 3) * dof + d_idx],
          u_position[b_idx * horizon * dof + (h_idx - 2) * dof + d_idx],
          u_position[b_idx * horizon * dof + (h_idx - 1) * dof + d_idx]);
        new_position     = u_clique.w;
        new_velocity     = (u_clique.w - u_clique.z) * dt; // 1 - 0
        new_acceleration = (u_clique.y - 2 * u_clique.z + u_clique.w) *
                           dt * dt;                        // 2 - 2.0 * 1 + 1
        new_jerk =
          (-1.0 * u_clique.x + 3 * u_clique.y - 3 * u_clique.z + u_clique.w) *
          dt * dt * dt;
      }

      // h_idx = h_idx + 1;
      out_position[b_idx * horizon * dof + h_idx * dof + d_idx] =
        new_position; // new_position;
      out_velocity[b_idx * horizon * dof + h_idx * dof + d_idx]     = new_velocity;
      out_acceleration[b_idx * horizon * dof + h_idx * dof + d_idx] =
        new_acceleration;
      out_jerk[b_idx * horizon * dof + h_idx * dof + d_idx] = new_jerk;
    }

    // Not used
    template<typename scalar_t>
    __global__ void position_clique_loop_coalesce_kernel(
      scalar_t *out_position_mem, scalar_t *out_velocity_mem,
      scalar_t *out_acceleration_mem, scalar_t *out_jerk_mem,
      const scalar_t *u_position, const scalar_t *start_position,
      const scalar_t *start_velocity, const scalar_t *start_acceleration,
      const scalar_t *traj_dt, const int batch_size, const int horizon,
      const int dof)
    {
      // data is stored as batch, dof, horizon
      // there are batch * horizon * dof threads:
      const int tid   = blockDim.x * blockIdx.x + threadIdx.x;
      const int b_idx = tid / (dof);
      const int d_idx = (tid - b_idx * dof);

      if ((b_idx >= batch_size) || (d_idx >= dof))
      {
        return;
      }

      const float dt =
        traj_dt[0]; // assume same dt across traj TODO: Implement variable dt

      // read start state:
      float u_arr[MAX_H];
      float out_pos[MAX_H], out_vel[MAX_H], out_acc[MAX_H], out_jerk[MAX_H];
      const int b_addrs = b_idx * horizon * dof;

#pragma unroll

      for (int i = 0; i < horizon; i++)
      {
        u_arr[i]    = u_position[b_addrs + d_idx * horizon + i];
        out_pos[i]  = 0;
        out_vel[i]  = 0;
        out_acc[i]  = 0;
        out_jerk[i] = 0;
      }

      out_pos[0]  = start_position[b_idx * dof + d_idx];
      out_vel[0]  = start_velocity[b_idx * dof + d_idx];
      out_acc[0]  = start_acceleration[b_idx * dof + d_idx];
      out_jerk[0] = 0.0;

      for (int h_idx = 1; h_idx < horizon; h_idx++)
      {
        // read actions: batch, horizon

        out_pos[h_idx] = u_arr[h_idx - 1];

        out_vel[h_idx] = (out_pos[h_idx] - out_pos[h_idx - 1]) * dt;  // 1 - 0
        out_acc[h_idx] =
          (out_vel[h_idx] - out_vel[h_idx - 1]) * dt;                 // 2 - 2.0 * 1 + 1
        out_jerk[h_idx] = (out_acc[h_idx] - out_acc[h_idx - 1]) * dt; // -1 3 -3 1
      }

      // write out:
      for (int h_idx = 0; h_idx < horizon; h_idx++)
      {
        out_position_mem[b_idx * horizon * dof + d_idx * horizon + h_idx] =
          out_pos[h_idx]; // new_position;
        out_velocity_mem[b_idx * horizon * dof + d_idx * horizon + h_idx] =
          out_vel[h_idx];
        out_acceleration_mem[b_idx * horizon * dof + d_idx * horizon + h_idx] =
          out_acc[h_idx];
        out_jerk_mem[b_idx * horizon * dof + d_idx * horizon + h_idx] =
          out_jerk[h_idx];
      }
    }

    // Not used
    template<typename scalar_t>
    __global__ void backward_position_clique_loop_coalesce_kernel(
      scalar_t *out_grad_position, const scalar_t *grad_position,
      const scalar_t *grad_velocity, const scalar_t *grad_acceleration,
      const scalar_t *grad_jerk, const scalar_t *traj_dt, const int batch_size,
      const int horizon, const int dof)
    {
      const int tid   = blockDim.x * blockIdx.x + threadIdx.x;
      const int b_idx = tid / (dof);
      const int d_idx = (tid - b_idx * dof);

      if ((b_idx >= batch_size) || (d_idx >= dof))
      {
        return;
      }
      const int b_addrs = b_idx * horizon * dof;

      // read gradients:
      float g_pos[MAX_H];
      float g_vel[MAX_H];
      float g_acc[MAX_H];
      float g_jerk[MAX_H];
      const float dt   = traj_dt[0];
      const float dt_2 = dt * dt;
      const float dt_3 = dt * dt * dt;
#pragma unroll

      for (int h_idx = 0; h_idx < horizon; h_idx++)
      {
        g_pos[h_idx]  = grad_position[b_addrs + d_idx * horizon + h_idx];
        g_vel[h_idx]  = grad_velocity[b_addrs + d_idx * horizon + h_idx];
        g_acc[h_idx]  = grad_acceleration[b_addrs + d_idx * horizon + h_idx];
        g_jerk[h_idx] = grad_jerk[b_addrs + d_idx * horizon + h_idx];
      }
#pragma unroll

      for (int i = 0; i < 4; i++)
      {
        g_vel[horizon + i]  = 0.0;
        g_acc[horizon + i]  = 0.0;
        g_jerk[horizon + i] = 0.0;
      }

      // compute gradient and sum
      for (int h_idx = 0; h_idx < horizon - 1; h_idx++)
      {
        g_pos[h_idx + 1] +=
          ((g_vel[h_idx + 1] - g_vel[h_idx + 2]) * dt +
           (g_acc[h_idx + 1] - 2 * g_acc[h_idx + 2] + g_acc[h_idx + 3]) * dt_2 +
           (1 * g_jerk[h_idx + 1] - 3 * g_jerk[h_idx + 2] +
            3 * g_jerk[h_idx + 3] - g_jerk[h_idx + 4]) *
           dt_3);
      }

      // write out:
      for (int h_idx = 0; h_idx < horizon - 1; h_idx++)
      {
        out_grad_position[b_addrs + d_idx * horizon + h_idx] = g_pos[h_idx + 1];
      }
      out_grad_position[b_addrs + d_idx * horizon + horizon - 1] = 0.0;
    }

    // Not used
    template<typename scalar_t>
    __global__ void backward_position_clique_kernel(
      scalar_t *out_grad_position, const scalar_t *grad_position,
      const scalar_t *grad_velocity, const scalar_t *grad_acceleration,
      const scalar_t *grad_jerk, const scalar_t *traj_dt, const int batch_size,
      const int horizon, const int dof)
    {
      // TODO: transpose h and dof to be able to directly read float2, float3, etc..
      const int tid   = blockDim.x * blockIdx.x + threadIdx.x;
      const int b_idx = tid / (horizon * dof);
      const int h_idx = (tid - b_idx * (horizon * dof)) / dof;
      const int d_idx = (tid - b_idx * horizon * dof - h_idx * dof);

      if ((b_idx >= batch_size) || (h_idx >= horizon) || (d_idx >= dof))
      {
        return;
      }
      const int b_addrs = b_idx * horizon * dof;

      if (h_idx == horizon - 1)
      {
        out_grad_position[b_addrs + (h_idx) * dof + d_idx] = 0.0;
        return;
      }

      // read gradients:
      const float dt = traj_dt[0];
      float g_u      = grad_position[b_addrs + (h_idx + 1) * dof + d_idx];

      float2 g_vel;
      float3 g_acc;
      float4 g_jerk;

      if (h_idx < horizon - 4) // && h_idx > 0)
      {
        g_vel = make_float2(grad_velocity[b_addrs + (h_idx + 1) * dof + d_idx],
                            grad_velocity[b_addrs + (h_idx + 2) * dof + d_idx]);
        g_acc = make_float3(grad_acceleration[b_addrs + (h_idx + 1) * dof + d_idx],
                            grad_acceleration[b_addrs + (h_idx + 2) * dof + d_idx],
                            grad_acceleration[b_addrs + (h_idx + 3) * dof + d_idx]);

        g_jerk = make_float4(grad_jerk[b_addrs + (h_idx + 1) * dof + d_idx],
                             grad_jerk[b_addrs + (h_idx + 2) * dof + d_idx],
                             grad_jerk[b_addrs + (h_idx + 3) * dof + d_idx],
                             grad_jerk[b_addrs + (h_idx + 4) * dof + d_idx]);
        g_u +=
          ((g_vel.x - g_vel.y) * dt +
           (g_acc.x - 2 * g_acc.y + g_acc.z) * dt * dt +
           (1 * g_jerk.x - 3 * g_jerk.y + 3 * g_jerk.z - g_jerk.w) * dt * dt * dt);
      }
      else if (h_idx == 0)
      {
        g_vel = make_float2(grad_velocity[b_addrs + (h_idx + 1) * dof + d_idx],
                            grad_velocity[b_addrs + (h_idx + 2) * dof + d_idx]);

        g_acc = make_float3(grad_acceleration[b_addrs + (h_idx + 1) * dof + d_idx],
                            grad_acceleration[b_addrs + (h_idx + 2) * dof + d_idx],
                            0.0);

        g_jerk = make_float4(grad_jerk[b_addrs + (h_idx + 1) * dof + d_idx],
                             grad_jerk[b_addrs + (h_idx + 2) * dof + d_idx],
                             grad_jerk[b_addrs + (h_idx + 3) * dof + d_idx], 0.0);
        g_u += ((g_vel.x - g_vel.y) * dt +
                (-1.0 * g_acc.x + 1 * g_acc.y) * dt * dt +
                (-1 * g_jerk.x + 2 * g_jerk.y - 1 * g_jerk.z) * dt * dt * dt);
      }
      else if (h_idx == horizon - 4)
      {
        g_vel = make_float2(grad_velocity[b_addrs + (h_idx + 1) * dof + d_idx],
                            grad_velocity[b_addrs + (h_idx + 2) * dof + d_idx]);
        g_acc = make_float3(grad_acceleration[b_addrs + (h_idx + 1) * dof + d_idx],
                            grad_acceleration[b_addrs + (h_idx + 2) * dof + d_idx],
                            grad_acceleration[b_addrs + (h_idx + 3) * dof + d_idx]);
        g_jerk = make_float4(grad_jerk[b_addrs + (h_idx + 1) * dof + d_idx],
                             grad_jerk[b_addrs + (h_idx + 2) * dof + d_idx],
                             grad_jerk[b_addrs + (h_idx + 3) * dof + d_idx], 0.0);
        g_u +=
          ((g_vel.x - g_vel.y) * dt +
           (g_acc.x - 2 * g_acc.y + g_acc.z) * dt * dt +
           (1 * g_jerk.x - 3 * g_jerk.y + 3 * g_jerk.z - g_jerk.w) * dt * dt * dt);
      }
      else if (h_idx == horizon - 3)
      {
        g_vel = make_float2(grad_velocity[b_addrs + (h_idx + 1) * dof + d_idx],
                            grad_velocity[b_addrs + (h_idx + 2) * dof + d_idx]);
        g_acc =
          make_float3(grad_acceleration[b_addrs + (h_idx + 1) * dof + d_idx],
                      grad_acceleration[b_addrs + (h_idx + 2) * dof + d_idx], 0);
        g_jerk =
          make_float4(grad_jerk[b_addrs + (h_idx + 1) * dof + d_idx],
                      grad_jerk[b_addrs + (h_idx + 2) * dof + d_idx], 0.0, 0.0);
        g_u +=
          ((g_vel.x - g_vel.y) * dt +
           (g_acc.x - 2 * g_acc.y + g_acc.z) * dt * dt +
           (1 * g_jerk.x - 3 * g_jerk.y + 3 * g_jerk.z - g_jerk.w) * dt * dt * dt);
      }
      else if (h_idx == horizon - 2)
      {
        g_vel =
          make_float2(grad_velocity[b_addrs + (h_idx + 1) * dof + d_idx], 0.0);
        g_acc = make_float3(grad_acceleration[b_addrs + (h_idx + 1) * dof + d_idx],
                            0, 0);
        g_jerk = make_float4(grad_jerk[b_addrs + (h_idx + 1) * dof + d_idx], 0.0,
                             0.0, 0.0);
        g_u +=
          ((g_vel.x - g_vel.y) * dt +
           (g_acc.x - 2 * g_acc.y + g_acc.z) * dt * dt +
           (1 * g_jerk.x - 3 * g_jerk.y + 3 * g_jerk.z - g_jerk.w) * dt * dt * dt);
      }

      out_grad_position[b_addrs + (h_idx) * dof + d_idx] = g_u;
    }
  } // namespace
}
std::vector<torch::Tensor>step_position_clique(
  torch::Tensor out_position, torch::Tensor out_velocity,
  torch::Tensor out_acceleration, torch::Tensor out_jerk,
  const torch::Tensor u_position, const torch::Tensor start_position,
  const torch::Tensor start_velocity, const torch::Tensor start_acceleration,
  const torch::Tensor traj_dt, const int batch_size, const int horizon,
  const int dof)
{
  using namespace Curobo::TensorStep;
  assert(horizon < MAX_H);

  const int k_size    = batch_size * dof;
  int threadsPerBlock = k_size;

  if (threadsPerBlock > 512)
  {
    threadsPerBlock = 512;
  }

  int blocksPerGrid   = (k_size + threadsPerBlock - 1) / threadsPerBlock;
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  AT_DISPATCH_FLOATING_TYPES(
    out_position.scalar_type(), "step_position_clique", ([&] {
    position_clique_loop_kernel<scalar_t>
      << < blocksPerGrid, threadsPerBlock, 0, stream >> > (
      out_position.data_ptr<scalar_t>(),
      out_velocity.data_ptr<scalar_t>(),
      out_acceleration.data_ptr<scalar_t>(),
      out_jerk.data_ptr<scalar_t>(), u_position.data_ptr<scalar_t>(),
      start_position.data_ptr<scalar_t>(),
      start_velocity.data_ptr<scalar_t>(),
      start_acceleration.data_ptr<scalar_t>(),
      traj_dt.data_ptr<scalar_t>(), batch_size, horizon, dof);
  }));

  C10_CUDA_KERNEL_LAUNCH_CHECK();

  return { out_position, out_velocity, out_acceleration, out_jerk };
}

std::vector<torch::Tensor>step_position_clique2(
  torch::Tensor out_position, torch::Tensor out_velocity,
  torch::Tensor out_acceleration, torch::Tensor out_jerk,
  const torch::Tensor u_position, const torch::Tensor start_position,
  const torch::Tensor start_velocity, const torch::Tensor start_acceleration,
  const torch::Tensor traj_dt, const int batch_size, const int horizon,
  const int dof,
  const int mode = -1)
{
  using namespace Curobo::TensorStep;

  assert(horizon > 5);

  // assert(horizon < MAX_H);
  const int k_size    = batch_size * dof * horizon;
  int threadsPerBlock = k_size;

  if (threadsPerBlock > 128)
  {
    threadsPerBlock = 128;
  }

  int blocksPerGrid   = (k_size + threadsPerBlock - 1) / threadsPerBlock;
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  if (mode == BWD_DIFF)
  {
    AT_DISPATCH_FLOATING_TYPES(
      out_position.scalar_type(), "step_position_clique", ([&] {
      position_clique_loop_kernel2<scalar_t, BWD_DIFF>
        << < blocksPerGrid, threadsPerBlock, 0, stream >> > (
        out_position.data_ptr<scalar_t>(),
        out_velocity.data_ptr<scalar_t>(),
        out_acceleration.data_ptr<scalar_t>(),
        out_jerk.data_ptr<scalar_t>(), u_position.data_ptr<scalar_t>(),
        start_position.data_ptr<scalar_t>(),
        start_velocity.data_ptr<scalar_t>(),
        start_acceleration.data_ptr<scalar_t>(),
        traj_dt.data_ptr<scalar_t>(), batch_size, horizon, dof);
    }));
  }
  else if (mode == CENTRAL_DIFF)
  {
    assert(u_position.sizes()[1] == horizon - 4);

    AT_DISPATCH_FLOATING_TYPES(
      out_position.scalar_type(), "step_position_clique", ([&] {
      position_clique_loop_kernel2<scalar_t, CENTRAL_DIFF>
        << < blocksPerGrid, threadsPerBlock, 0, stream >> > (
        out_position.data_ptr<scalar_t>(),
        out_velocity.data_ptr<scalar_t>(),
        out_acceleration.data_ptr<scalar_t>(),
        out_jerk.data_ptr<scalar_t>(), u_position.data_ptr<scalar_t>(),
        start_position.data_ptr<scalar_t>(),
        start_velocity.data_ptr<scalar_t>(),
        start_acceleration.data_ptr<scalar_t>(),
        traj_dt.data_ptr<scalar_t>(), batch_size, horizon, dof);
    }));
  }
  else
  {
    assert(false);
  }
  C10_CUDA_KERNEL_LAUNCH_CHECK();

  return { out_position, out_velocity, out_acceleration, out_jerk };
}

std::vector<torch::Tensor>step_position_clique2_idx(
  torch::Tensor out_position, torch::Tensor out_velocity,
  torch::Tensor out_acceleration, torch::Tensor out_jerk,
  const torch::Tensor u_position, const torch::Tensor start_position,
  const torch::Tensor start_velocity, const torch::Tensor start_acceleration,
  const torch::Tensor start_idx, const torch::Tensor traj_dt,
  const int batch_size, const int horizon, const int dof,
  const int mode = -1)
{
  using namespace Curobo::TensorStep;

  // assert(horizon < MAX_H);
  assert(horizon > 5);


  const int k_size    = batch_size * dof * horizon;
  int threadsPerBlock = k_size;

  if (threadsPerBlock > 128)
  {
    threadsPerBlock = 128;
  }

  int blocksPerGrid   = (k_size + threadsPerBlock - 1) / threadsPerBlock;
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  if (mode == BWD_DIFF)
  {
    assert(false);
    AT_DISPATCH_FLOATING_TYPES(
      out_position.scalar_type(), "step_position_clique", ([&] {
      position_clique_loop_idx_kernel2<scalar_t, BWD_DIFF>
        << < blocksPerGrid, threadsPerBlock, 0, stream >> > (
        out_position.data_ptr<scalar_t>(),
        out_velocity.data_ptr<scalar_t>(),
        out_acceleration.data_ptr<scalar_t>(),
        out_jerk.data_ptr<scalar_t>(), u_position.data_ptr<scalar_t>(),
        start_position.data_ptr<scalar_t>(),
        start_velocity.data_ptr<scalar_t>(),
        start_acceleration.data_ptr<scalar_t>(),
        start_idx.data_ptr<int32_t>(), traj_dt.data_ptr<scalar_t>(),
        batch_size, horizon, dof);
    }));
  }

  else if (mode == CENTRAL_DIFF)
  {
    assert(u_position.sizes()[1] == horizon - 4);

    AT_DISPATCH_FLOATING_TYPES(
      out_position.scalar_type(), "step_position_clique", ([&] {
      position_clique_loop_idx_kernel2<scalar_t, CENTRAL_DIFF>
        << < blocksPerGrid, threadsPerBlock, 0, stream >> > (
        out_position.data_ptr<scalar_t>(),
        out_velocity.data_ptr<scalar_t>(),
        out_acceleration.data_ptr<scalar_t>(),
        out_jerk.data_ptr<scalar_t>(), u_position.data_ptr<scalar_t>(),
        start_position.data_ptr<scalar_t>(),
        start_velocity.data_ptr<scalar_t>(),
        start_acceleration.data_ptr<scalar_t>(),
        start_idx.data_ptr<int32_t>(), traj_dt.data_ptr<scalar_t>(),
        batch_size, horizon, dof);
    }));
  }
  else
  {
    assert(false);
  }
  C10_CUDA_KERNEL_LAUNCH_CHECK();

  return { out_position, out_velocity, out_acceleration, out_jerk };
}

std::vector<torch::Tensor>backward_step_position_clique(
  torch::Tensor out_grad_position, const torch::Tensor grad_position,
  const torch::Tensor grad_velocity, const torch::Tensor grad_acceleration,
  const torch::Tensor grad_jerk, const torch::Tensor traj_dt,
  const int batch_size, const int horizon, const int dof)
{
  using namespace Curobo::TensorStep;

  assert(horizon < MAX_H - 4);
  const int k_size    = batch_size * dof;
  int threadsPerBlock = k_size;

  if (threadsPerBlock > 128)
  {
    threadsPerBlock = 128;
  }

  int blocksPerGrid   = (k_size + threadsPerBlock - 1) / threadsPerBlock;
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  AT_DISPATCH_FLOATING_TYPES(
    out_grad_position.scalar_type(), "backward_step_position_clique", ([&] {
    backward_position_clique_loop_kernel<scalar_t>
      << < blocksPerGrid, threadsPerBlock, 0, stream >> > (
      out_grad_position.data_ptr<scalar_t>(),
      grad_position.data_ptr<scalar_t>(),
      grad_velocity.data_ptr<scalar_t>(),
      grad_acceleration.data_ptr<scalar_t>(),
      grad_jerk.data_ptr<scalar_t>(), traj_dt.data_ptr<scalar_t>(),
      batch_size, horizon, dof);
  }));

  C10_CUDA_KERNEL_LAUNCH_CHECK();

  return { out_grad_position };
}

std::vector<torch::Tensor>backward_step_position_clique2(
  torch::Tensor out_grad_position, const torch::Tensor grad_position,
  const torch::Tensor grad_velocity, const torch::Tensor grad_acceleration,
  const torch::Tensor grad_jerk, const torch::Tensor traj_dt,
  const int batch_size, const int horizon, const int dof,
  const int mode = -1)
{
  // assert(horizon < MAX_H - 4);
  using namespace Curobo::TensorStep;

  assert(horizon > 5);


  // const int k_size = batch_size * dof;
  const int k_size    = batch_size * dof * horizon;
  int threadsPerBlock = k_size;

  if (threadsPerBlock > 128)
  {
    threadsPerBlock = 128;
  }

  int blocksPerGrid   = (k_size + threadsPerBlock - 1) / threadsPerBlock;
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  if (mode == BWD_DIFF)
  {
    assert(false); // not supported anymore
    AT_DISPATCH_FLOATING_TYPES(
      out_grad_position.scalar_type(), "backward_step_position_clique", ([&] {
      backward_position_clique_loop_backward_difference_kernel2<scalar_t>
        << < blocksPerGrid, threadsPerBlock, 0, stream >> > (
        out_grad_position.data_ptr<scalar_t>(),
        grad_position.data_ptr<scalar_t>(),
        grad_velocity.data_ptr<scalar_t>(),
        grad_acceleration.data_ptr<scalar_t>(),
        grad_jerk.data_ptr<scalar_t>(), traj_dt.data_ptr<scalar_t>(),
        batch_size, horizon, dof);
    }));
  }
  else if (mode == CENTRAL_DIFF)
  {
    assert(out_grad_position.sizes()[1] == horizon - 4);
    AT_DISPATCH_FLOATING_TYPES(
      out_grad_position.scalar_type(), "backward_step_position_clique", ([&] {
      backward_position_clique_loop_central_difference_kernel2<scalar_t>
        << < blocksPerGrid, threadsPerBlock, 0, stream >> > (
        out_grad_position.data_ptr<scalar_t>(),
        grad_position.data_ptr<scalar_t>(),
        grad_velocity.data_ptr<scalar_t>(),
        grad_acceleration.data_ptr<scalar_t>(),
        grad_jerk.data_ptr<scalar_t>(), traj_dt.data_ptr<scalar_t>(),
        batch_size, horizon, dof);
    }));
  }
  else
  {
    assert(false);
  }

  C10_CUDA_KERNEL_LAUNCH_CHECK();

  return { out_grad_position };
}

std::vector<torch::Tensor>
step_acceleration(torch::Tensor out_position, torch::Tensor out_velocity,
                  torch::Tensor out_acceleration, torch::Tensor out_jerk,
                  const torch::Tensor u_acc, const torch::Tensor start_position,
                  const torch::Tensor start_velocity,
                  const torch::Tensor start_acceleration,
                  const torch::Tensor traj_dt, const int batch_size,
                  const int horizon, const int dof, const bool use_rk2 = true)
{
  assert(horizon < MAX_H);
  using namespace Curobo::TensorStep;

  const int k_size    = batch_size * dof;
  int threadsPerBlock = k_size;

  if (threadsPerBlock > 512)
  {
    threadsPerBlock = 512;
  }

  int blocksPerGrid   = (k_size + threadsPerBlock - 1) / threadsPerBlock;
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  if (use_rk2)
  {
    AT_DISPATCH_FLOATING_TYPES(
      out_position.scalar_type(), "step_acceleration", ([&] {
      acceleration_loop_rk2_kernel<scalar_t>
        << < blocksPerGrid, threadsPerBlock, 0, stream >> > (
        out_position.data_ptr<scalar_t>(),
        out_velocity.data_ptr<scalar_t>(),
        out_acceleration.data_ptr<scalar_t>(),
        out_jerk.data_ptr<scalar_t>(), u_acc.data_ptr<scalar_t>(),
        start_position.data_ptr<scalar_t>(),
        start_velocity.data_ptr<scalar_t>(),
        start_acceleration.data_ptr<scalar_t>(),
        traj_dt.data_ptr<scalar_t>(), batch_size, horizon, dof);
    }));
  }

  else
  {
    AT_DISPATCH_FLOATING_TYPES(
      out_position.scalar_type(), "step_acceleration", ([&] {
      acceleration_loop_kernel<scalar_t>
        << < blocksPerGrid, threadsPerBlock, 0, stream >> > (
        out_position.data_ptr<scalar_t>(),
        out_velocity.data_ptr<scalar_t>(),
        out_acceleration.data_ptr<scalar_t>(),
        out_jerk.data_ptr<scalar_t>(), u_acc.data_ptr<scalar_t>(),
        start_position.data_ptr<scalar_t>(),
        start_velocity.data_ptr<scalar_t>(),
        start_acceleration.data_ptr<scalar_t>(),
        traj_dt.data_ptr<scalar_t>(), batch_size, horizon, dof);
    }));
  }
  C10_CUDA_KERNEL_LAUNCH_CHECK();

  return { out_position, out_velocity, out_acceleration, out_jerk };
}

std::vector<torch::Tensor>step_acceleration_idx(
  torch::Tensor out_position, torch::Tensor out_velocity,
  torch::Tensor out_acceleration, torch::Tensor out_jerk,
  const torch::Tensor u_acc, const torch::Tensor start_position,
  const torch::Tensor start_velocity, const torch::Tensor start_acceleration,
  const torch::Tensor start_idx, const torch::Tensor traj_dt,
  const int batch_size, const int horizon, const int dof,
  const bool use_rk2 = true)
{
  assert(horizon < MAX_H);
  using namespace Curobo::TensorStep;


  const int k_size    = batch_size * dof;
  int threadsPerBlock = k_size;

  if (threadsPerBlock > 512)
  {
    threadsPerBlock = 512;
  }

  int blocksPerGrid   = (k_size + threadsPerBlock - 1) / threadsPerBlock;
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  if (use_rk2)
  {
    AT_DISPATCH_FLOATING_TYPES(
      out_position.scalar_type(), "step_acceleration", ([&] {
      acceleration_loop_idx_rk2_kernel<scalar_t>
        << < blocksPerGrid, threadsPerBlock, 0, stream >> > (
        out_position.data_ptr<scalar_t>(),
        out_velocity.data_ptr<scalar_t>(),
        out_acceleration.data_ptr<scalar_t>(),
        out_jerk.data_ptr<scalar_t>(), u_acc.data_ptr<scalar_t>(),
        start_position.data_ptr<scalar_t>(),
        start_velocity.data_ptr<scalar_t>(),
        start_acceleration.data_ptr<scalar_t>(),
        start_idx.data_ptr<int32_t>(), traj_dt.data_ptr<scalar_t>(),
        batch_size, horizon, dof);
    }));
  }
  else
  {
    AT_DISPATCH_FLOATING_TYPES(
      out_position.scalar_type(), "step_acceleration", ([&] {
      acceleration_loop_idx_kernel<scalar_t>
        << < blocksPerGrid, threadsPerBlock, 0, stream >> > (
        out_position.data_ptr<scalar_t>(),
        out_velocity.data_ptr<scalar_t>(),
        out_acceleration.data_ptr<scalar_t>(),
        out_jerk.data_ptr<scalar_t>(), u_acc.data_ptr<scalar_t>(),
        start_position.data_ptr<scalar_t>(),
        start_velocity.data_ptr<scalar_t>(),
        start_acceleration.data_ptr<scalar_t>(),
        start_idx.data_ptr<int32_t>(), traj_dt.data_ptr<scalar_t>(),
        batch_size, horizon, dof);
    }));
  }
  C10_CUDA_KERNEL_LAUNCH_CHECK();

  return { out_position, out_velocity, out_acceleration, out_jerk };
}
