Add torch pendulum (#32)

Author: astomodynamics

CMakeLists.txt

@@ -30,7 +30,7 @@ option(CDDP-CPP_BUILD_TESTS "Whether to build tests." ON)
 option(CDDP-CPP_GUROBI "Whether to use Gurobi solver." OFF)
 option(GUROBI_ROOT "Path to Gurobi installation" "") 
 set(GUROBI_ROOT /home/tom/.local/lib/gurobi1103/linux64)
-option(CDDP-CPP_TORCH "Whether to use LibTorch." ON)
+set(CDDP-CPP_TORCH "Whether to use LibTorch." ON) # cannot be turned off
 option(CDDP-CPP_TORCH_GPU "Whether to use GPU." ON)
 
 # Find packages
@@ -66,6 +66,55 @@ if (CDDP-CPP_BUILD_TESTS)
   include(GoogleTest)
 endif()
 
+# LibTorch
+if (CDDP-CPP_TORCH)
+  if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/build/libtorch/share/cmake/Torch/TorchConfig.cmake)
+    message(STATUS "Found LibTorch at ${CMAKE_CURRENT_SOURCE_DIR}/build/libtorch")
+  else()
+    message(STATUS "Downloading LibTorch...")
+    # Download and extract LibTorch
+    if (CDDP-CPP_TORCH_GPU)
+      set(LIBTORCH_URL "https://download.pytorch.org/libtorch/cu124/libtorch-cxx11-abi-shared-with-deps-2.5.1%2Bcu124.zip")
+    else()
+      set(LIBTORCH_URL "https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.5.1%2Bcpu.zip")
+    endif()
+
+    # Set the download directory
+    set(DOWLOAD_DIR ${CMAKE_CURRENT_SOURCE_DIR}/build)
+
+    # Create the download directory
+    file(MAKE_DIRECTORY ${DOWLOAD_DIR})
+
+    # Download the file 
+    file(DOWNLOAD ${LIBTORCH_URL} ${DOWLOAD_DIR}/libtorch-shared-with-deps-latest.zip
+      STATUS DOWNLOAD_STATUS
+      SHOW_PROGRESS
+    )
+
+    # Check if the download was successful
+    list(GET DOWNLOAD_STATUS 0 DOWNLOAD_STATUS_CODE)
+    if(NOT DOWNLOAD_STATUS_CODE EQUAL 0)
+      message(FATAL_ERROR "Failed to download LibTorch.")
+    endif()
+    
+    # Extract the file    
+    execute_process(
+      COMMAND ${CMAKE_COMMAND} -E tar xvf ${DOWLOAD_DIR}/libtorch-shared-with-deps-latest.zip
+      WORKING_DIRECTORY ${DOWLOAD_DIR}
+    )
+
+    # Remove the zip file
+    file(REMOVE ${DOWLOAD_DIR}/libtorch-shared-with-deps-latest.zip)
+  endif()
+
+  # Set the path to the LibTorch installation
+  set(LIBTORCH_DIR ${CMAKE_CURRENT_SOURCE_DIR}/build/libtorch)
+
+  find_package(Torch REQUIRED PATHS ${LIBTORCH_DIR} NO_DEFAULT_PATH)
+  message(STATUS "Found LibTorch: ${TORCH_LIBRARIES}")
+endif()
+
+
 # Include directories
 include_directories(
   ${CMAKE_CURRENT_SOURCE_DIR}/include
@@ -79,6 +128,7 @@ set(cddp_core_srcs
   src/cddp_core/objective.cpp
   src/cddp_core/constraint.cpp
   src/cddp_core/cddp_core.cpp
+  src/cddp_core/torch_dynamical_system.cpp
 )
 
 set(dynamics_model_srcs
@@ -102,11 +152,13 @@ target_link_libraries(${PROJECT_NAME}
   Python3::Python
   Python3::Module
   Python3::NumPy
+  ${TORCH_LIBRARIES}
 ) 
 
 target_include_directories(${PROJECT_NAME} PUBLIC
   $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include/cddp-cpp>
   $<INSTALL_INTERFACE:include> 
+  ${TORCH_INCLUDE_DIRS}
 )
 
 # Gurobi 
@@ -132,57 +184,6 @@ if (CDDP-CPP_GUROBI)
   endif()
 endif()
 
-# LibTorch
-if (CDDP-CPP_TORCH)
-  if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/build/libtorch/share/cmake/Torch/TorchConfig.cmake)
-    message(STATUS "Found LibTorch at ${CMAKE_CURRENT_SOURCE_DIR}/build/libtorch")
-  else()
-    message(STATUS "Downloading LibTorch...")
-    # Download and extract LibTorch
-    if (CDDP-CPP_TORCH_GPU)
-      set(LIBTORCH_URL "https://download.pytorch.org/libtorch/cu124/libtorch-cxx11-abi-shared-with-deps-2.5.1%2Bcu124.zip")
-    else()
-      set(LIBTORCH_URL "https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.5.1%2Bcpu.zip")
-    endif()
-
-    # Set the download directory
-    set(DOWLOAD_DIR ${CMAKE_CURRENT_SOURCE_DIR}/build)
-
-    # Create the download directory
-    file(MAKE_DIRECTORY ${DOWLOAD_DIR})
-
-    # Download the file 
-    file(DOWNLOAD ${LIBTORCH_URL} ${DOWLOAD_DIR}/libtorch-shared-with-deps-latest.zip
-      STATUS DOWNLOAD_STATUS
-      SHOW_PROGRESS
-    )
-
-    # Check if the download was successful
-    list(GET DOWNLOAD_STATUS 0 DOWNLOAD_STATUS_CODE)
-    if(NOT DOWNLOAD_STATUS_CODE EQUAL 0)
-      message(FATAL_ERROR "Failed to download LibTorch.")
-    endif()
-    
-    # Extract the file    
-    execute_process(
-      COMMAND ${CMAKE_COMMAND} -E tar xvf ${DOWLOAD_DIR}/libtorch-shared-with-deps-latest.zip
-      WORKING_DIRECTORY ${DOWLOAD_DIR}
-    )
-
-    # Remove the zip file
-    file(REMOVE ${DOWLOAD_DIR}/libtorch-shared-with-deps-latest.zip)
-  endif()
-
-  # Set the path to the LibTorch installation
-  set(LIBTORCH_DIR ${CMAKE_CURRENT_SOURCE_DIR}/build/libtorch)
-
-  find_package(Torch REQUIRED PATHS ${LIBTORCH_DIR} NO_DEFAULT_PATH)
-  target_link_libraries(${PROJECT_NAME} ${TORCH_LIBRARIES})
-  target_include_directories(${PROJECT_NAME} PUBLIC ${TORCH_INCLUDE_DIRS})
-  message(STATUS "Found LibTorch: ${TORCH_LIBRARIES}")
-endif()
-
-
 # Build and register tests.
 if (CDDP-CPP_BUILD_TESTS)
   add_subdirectory(tests) 

include/cddp-cpp/cddp.hpp

@@ -21,13 +21,14 @@
 #include <vector>
 #include <Eigen/Dense>
 
-
-
+// #include "cddp-cpp/sdqp.hpp"
 #include "cddp_core/dynamical_system.hpp"
 #include "cddp_core/objective.hpp"
 #include "cddp_core/constraint.hpp"
 #include "cddp_core/cddp_core.hpp"
 
+#include "cddp_core/torch_dynamical_system.hpp"
+
 // Models
 #include "dynamics_model/pendulum.hpp"
 #include "dynamics_model/dubins_car.hpp"

include/cddp-cpp/cddp_core/torch_dynamical_system.hpp

@@ -0,0 +1,74 @@
+/*
+ Copyright 2024 Tomo Sasaki
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+      https://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+*/
+#ifndef CDDP_TORCH_DYNAMICAL_SYSTEM_HPP
+#define CDDP_TORCH_DYNAMICAL_SYSTEM_HPP
+
+#include "cddp_core/dynamical_system.hpp" 
+#include <torch/torch.h>
+#include <Eigen/Dense>
+
+namespace cddp {
+
+class DynamicsModelInterface : public torch::nn::Module {
+public:
+    virtual torch::Tensor forward(std::vector<torch::Tensor> inputs) = 0;
+    virtual ~DynamicsModelInterface() = default;
+};
+
+class TorchDynamicalSystem : public DynamicalSystem {
+public:
+    TorchDynamicalSystem(int state_dim, 
+                        int control_dim, 
+                        double timestep,
+                        std::string integration_type,
+                        std::shared_ptr<DynamicsModelInterface> model,
+                        bool use_gpu = false);
+
+    // Override core dynamics methods
+    Eigen::VectorXd getContinuousDynamics(const Eigen::VectorXd& state, 
+                                         const Eigen::VectorXd& control) const override;
+
+    Eigen::MatrixXd getStateJacobian(const Eigen::VectorXd& state, 
+                                    const Eigen::VectorXd& control) const override;
+
+    Eigen::MatrixXd getControlJacobian(const Eigen::VectorXd& state, 
+                                      const Eigen::VectorXd& control) const override;
+                                      
+    // Add batch processing capability
+    std::vector<Eigen::VectorXd> getBatchDynamics(
+        const std::vector<Eigen::VectorXd>& states,
+        const std::vector<Eigen::VectorXd>& controls) const;
+
+    // Optional: Override Hessian computations if needed
+    Eigen::MatrixXd getStateHessian(const Eigen::VectorXd& state,
+                                   const Eigen::VectorXd& control) const override;
+    
+    Eigen::MatrixXd getControlHessian(const Eigen::VectorXd& state,
+                                     const Eigen::VectorXd& control) const override;
+
+private:
+    // Helper methods for tensor conversions
+    torch::Tensor eigenToTorch(const Eigen::VectorXd& eigen_vec, bool requires_grad = false) const;
+    Eigen::VectorXd torchToEigen(const torch::Tensor& tensor) const;
+
+    std::shared_ptr<DynamicsModelInterface> model_;
+    bool use_gpu_;
+    torch::Device device_;
+};
+
+} // namespace cddp
+
+#endif // CDDP_TORCH_DYNAMICAL_SYSTEM_HPP