From 8be9ecb3c3ef46125d5677fcd17623c504784927 Mon Sep 17 00:00:00 2001 From: Tomo Date: Tue, 24 Dec 2024 17:59:08 -0500 Subject: [PATCH] Neural dynamics update (#61) * Re-create neural dynamics and tested --- .gitignore | 3 +- CMakeLists.txt | 4 +- include/cddp-cpp/cddp.hpp | 3 +- include/cddp-cpp/cddp_core/cddp_core.hpp | 8 +- .../cddp_core/neural_dynamical_system.hpp | 139 +++++++ .../cddp_core/torch_dynamical_system.hpp | 74 ---- src/cddp_core/cddp_core.cpp | 1 + src/cddp_core/neural_dynamical_system.cpp | 202 ++++++++++ src/cddp_core/torch_dynamical_system.cpp | 182 --------- tests/CMakeLists.txt | 22 +- tests/dynamics_model/test_neural_cartpole.cpp | 0 tests/dynamics_model/test_neural_pendulum.cpp | 299 +++++++++++++++ .../dynamics_model/test_neural_quadrotor.cpp | 0 tests/dynamics_model/test_torch_pendulum.cpp | 346 ------------------ 14 files changed, 661 insertions(+), 622 deletions(-) create mode 100644 include/cddp-cpp/cddp_core/neural_dynamical_system.hpp delete mode 100644 include/cddp-cpp/cddp_core/torch_dynamical_system.hpp create mode 100644 src/cddp_core/neural_dynamical_system.cpp delete mode 100644 src/cddp_core/torch_dynamical_system.cpp create mode 100644 tests/dynamics_model/test_neural_cartpole.cpp create mode 100644 tests/dynamics_model/test_neural_pendulum.cpp create mode 100644 tests/dynamics_model/test_neural_quadrotor.cpp delete mode 100644 tests/dynamics_model/test_torch_pendulum.cpp diff --git a/.gitignore b/.gitignore index 9c05b2f..19c0e13 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ .vscode/ build/ plots/ -results/frames \ No newline at end of file +results/frames +models/ \ No newline at end of file diff --git a/CMakeLists.txt b/CMakeLists.txt index 42beca5..ba18d38 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -40,7 +40,7 @@ option(GUROBI_ROOT "Path to Gurobi installation" "") set(GUROBI_ROOT $ENV{HOME}/.local/lib/gurobi1103/linux64) # LibTorch Configuration -option(CDDP_CPP_TORCH "Whether to use LibTorch" ON) +set(CDDP_CPP_TORCH "Whether to use LibTorch" ON) option(CDDP_CPP_TORCH_GPU "Whether to use GPU support in LibTorch" ON) set(LIBTORCH_DIR $ENV{HOME}/.local/lib/libtorch CACHE PATH "Path to local LibTorch installation") # FIXME: Change this to your local LibTorch installation directory @@ -198,7 +198,7 @@ set(cddp_core_srcs src/cddp_core/cddp_core.cpp src/cddp_core/asddp_core.cpp src/cddp_core/logddp_core.cpp - src/cddp_core/torch_dynamical_system.cpp + src/cddp_core/neural_dynamical_system.cpp ) set(dynamics_model_srcs diff --git a/include/cddp-cpp/cddp.hpp b/include/cddp-cpp/cddp.hpp index cf4a7d2..34eeb01 100644 --- a/include/cddp-cpp/cddp.hpp +++ b/include/cddp-cpp/cddp.hpp @@ -29,8 +29,7 @@ #include "cddp_core/helper.hpp" #include "cddp_core/boxqp.hpp" #include "cddp_core/qp_solver.hpp" - -#include "cddp_core/torch_dynamical_system.hpp" +#include "cddp_core/neural_dynamical_system.hpp" // Models #include "dynamics_model/pendulum.hpp" diff --git a/include/cddp-cpp/cddp_core/cddp_core.hpp b/include/cddp-cpp/cddp_core/cddp_core.hpp index 90f8526..a8a6c3c 100644 --- a/include/cddp-cpp/cddp_core/cddp_core.hpp +++ b/include/cddp-cpp/cddp_core/cddp_core.hpp @@ -258,13 +258,13 @@ class CDDP { private: bool initialized_ = false; // Initialization flag + // Problem Data - std::unique_ptr system_; // Eigen-based dynamical system - // std::unique_ptr torch_system_; // Torch-based dynamical system (learned dynamics model) - std::unique_ptr objective_; // Objective function + std::unique_ptr system_; + std::unique_ptr objective_; std::map> constraint_set_; std::unique_ptr log_barrier_; - Eigen::VectorXd initial_state_; // Initial state of the system + Eigen::VectorXd initial_state_; Eigen::VectorXd reference_state_; // Desired reference state int horizon_; // Time horizon for the problem double timestep_; // Time step for the problem diff --git a/include/cddp-cpp/cddp_core/neural_dynamical_system.hpp b/include/cddp-cpp/cddp_core/neural_dynamical_system.hpp new file mode 100644 index 0000000..0ef1d30 --- /dev/null +++ b/include/cddp-cpp/cddp_core/neural_dynamical_system.hpp @@ -0,0 +1,139 @@ +/* + 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_NEURAL_DYNAMICAL_SYSTEM_HPP +#define CDDP_NEURAL_DYNAMICAL_SYSTEM_HPP + +#include +#include +#include +#include +#include "cddp_core/dynamical_system.hpp" + +namespace cddp { +/** + * @brief Interface for a neural network model representing system dynamics. + */ + +class DynamicsModelInterface : public torch::nn::Module { +public: + virtual torch::Tensor forward(std::vector inputs) = 0; + virtual ~DynamicsModelInterface() = default; +}; + +/** + * @brief A NeuralDynamicalSystem that uses a PyTorch model to represent system dynamics. + */ +class NeuralDynamicalSystem : public DynamicalSystem { +public: + /** + * @brief Construct a new NeuralDynamicalSystem object + * + * @param state_dim Dimension of the system state + * @param control_dim Dimension of the system control + * @param timestep Integration timestep + * @param integration_type Type of numerical integration (Euler, Heun, RK3, RK4) + * @param model A torch::nn::Module (e.g. an MLP) representing the learned dynamics, + * @param device Device to run the model on (CPU or CUDA) + */ + NeuralDynamicalSystem(int state_dim, + int control_dim, + double timestep, + const std::string& integration_type, + std::shared_ptr model, + torch::Device device = torch::kCPU); + + /** + * @brief Compute continuous-time dynamics: x_dot = f(x, u). + * + * @param state Current state (Eigen vector) + * @param control Current control (Eigen vector) + * @return Eigen::VectorXd Continuous-time derivative of state + */ + Eigen::VectorXd getContinuousDynamics(const Eigen::VectorXd& state, + const Eigen::VectorXd& control) const override; + + /** + * @brief Compute discrete-time dynamics: x_{t+1} = f(x_t, u_t). + * + * @param state Current state (Eigen vector) + * @param control Current control (Eigen vector) + * @return Eigen::VectorXd Discrete next state + */ + Eigen::VectorXd getDiscreteDynamics(const Eigen::VectorXd& state, + const Eigen::VectorXd& control) const override; + + /** + * @brief Jacobian of the dynamics w.r.t. state: df/dx + * + * @param state Current state (Eigen vector) + * @param control Current control (Eigen vector) + * @return Eigen::MatrixXd Jacobian df/dx + */ + Eigen::MatrixXd getStateJacobian(const Eigen::VectorXd& state, + const Eigen::VectorXd& control) const override; + + /** + * @brief Jacobian of the dynamics w.r.t. control: df/du + * + * @param state Current state (Eigen vector) + * @param control Current control (Eigen vector) + * @return Eigen::MatrixXd Jacobian df/du + */ + Eigen::MatrixXd getControlJacobian(const Eigen::VectorXd& state, + const Eigen::VectorXd& control) const override; + + /** + * @brief Hessian of the dynamics w.r.t. state (flattened or block representation). + * + * @param state Current state (Eigen vector) + * @param control Current control (Eigen vector) + * @return Eigen::MatrixXd Hessian d^2f/dx^2 + */ + Eigen::MatrixXd getStateHessian(const Eigen::VectorXd& state, + const Eigen::VectorXd& control) const override; + + /** + * @brief Hessian of the dynamics w.r.t. control (flattened or block representation). + * + * @param state Current state (Eigen vector) + * @param control Current control (Eigen vector) + * @return Eigen::MatrixXd Hessian d^2f/du^2 + */ + Eigen::MatrixXd getControlHessian(const Eigen::VectorXd& state, + const Eigen::VectorXd& control) const override; + + /** + * @brief Hessian of the dynamics w.r.t. state and control (flattened or block representation). + * + * @param state Current state (Eigen vector) + * @param control Current control (Eigen vector) + * @return Eigen::MatrixXd Hessian d^2f/dudx + */ + Eigen::MatrixXd getCrossHessian(const Eigen::VectorXd& state, + const Eigen::VectorXd& control) const override; + +private: + std::shared_ptr model_; + torch::Device device_; + + // 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; +}; +} // namespace cddp + +#endif // CDDP_NEURAL_DYNAMICAL_SYSTEM_HPP diff --git a/include/cddp-cpp/cddp_core/torch_dynamical_system.hpp b/include/cddp-cpp/cddp_core/torch_dynamical_system.hpp deleted file mode 100644 index d65d100..0000000 --- a/include/cddp-cpp/cddp_core/torch_dynamical_system.hpp +++ /dev/null @@ -1,74 +0,0 @@ -/* - 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 -#include - -namespace cddp { - -class DynamicsModelInterface : public torch::nn::Module { -public: - virtual torch::Tensor forward(std::vector 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 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 getBatchDynamics( - const std::vector& states, - const std::vector& 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 model_; - bool use_gpu_; - torch::Device device_; -}; - -} // namespace cddp - -#endif // CDDP_TORCH_DYNAMICAL_SYSTEM_HPP \ No newline at end of file diff --git a/src/cddp_core/cddp_core.cpp b/src/cddp_core/cddp_core.cpp index 6d371dd..7e954d9 100644 --- a/src/cddp_core/cddp_core.cpp +++ b/src/cddp_core/cddp_core.cpp @@ -442,6 +442,7 @@ bool CDDP::solveBackwardPass() { const Eigen::VectorXd& x = X_[t]; const Eigen::VectorXd& u = U_[t]; + // TODO: Precompute Jacobians and store them? // Get continuous dynamics Jacobians const auto [Fx, Fu] = system_->getJacobians(x, u); diff --git a/src/cddp_core/neural_dynamical_system.cpp b/src/cddp_core/neural_dynamical_system.cpp new file mode 100644 index 0000000..1a5a47b --- /dev/null +++ b/src/cddp_core/neural_dynamical_system.cpp @@ -0,0 +1,202 @@ +/* + 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. +*/ + +#include +#include +#include "cddp_core/neural_dynamical_system.hpp" + +namespace cddp { + +NeuralDynamicalSystem::NeuralDynamicalSystem(int state_dim, + int control_dim, + double timestep, + const std::string& integration_type, + std::shared_ptr model, + torch::Device device) + : DynamicalSystem(state_dim, control_dim, timestep, integration_type) + , model_(std::move(model)) + , device_(device) +{ + if (!model_) { + throw std::runtime_error("NeuralDynamicalSystem: Received null model pointer."); + } + + // Move model to desired device + model_->to(device_); + + // Check model dimension + auto dummy_state = torch::zeros({1, state_dim_}, torch::kDouble).to(device_); + auto dummy_control = torch::zeros({1, control_dim_}, torch::kDouble).to(device_); + auto output = model_->forward({dummy_state, dummy_control}); + + // Expected output to be shape [1, state_dim_] + if (output.dim() != 2 || output.size(0) != 1 || output.size(1) != state_dim_) { + throw std::runtime_error( + "NeuralDynamicalSystem: Model output shape mismatch. " + "Expected [1, " + std::to_string(state_dim_) + "] but got " + + std::to_string(output.sizes()[0]) + " x " + std::to_string(output.sizes()[1])); + } +} + +// ---------------------------------------------------------------------------- +// getContinuousDynamics +// ---------------------------------------------------------------------------- +Eigen::VectorXd NeuralDynamicalSystem::getContinuousDynamics(const Eigen::VectorXd& state, + const Eigen::VectorXd& control) const +{ + // We'll assume the model produces x_dot = f(x,u). + // Or if it produces next state, adapt accordingly. + torch::NoGradGuard no_grad; // We don't need gradient for forward pass + + auto state_tensor = eigenToTorch(state); // shape: [1, state_dim_] + auto control_tensor = eigenToTorch(control); // shape: [1, control_dim_] + + auto output = model_->forward({state_tensor, control_tensor}); + // Convert to Eigen + return torchToEigen(output); +} + +// ---------------------------------------------------------------------------- +// getDiscreteDynamics +// ---------------------------------------------------------------------------- +Eigen::VectorXd NeuralDynamicalSystem::getDiscreteDynamics(const Eigen::VectorXd& state, + const Eigen::VectorXd& control) const +{ + // A typical approach: x_next = x + x_dot*dt if the model outputs x_dot + // If your model directly outputs next-state, just return the model result. + // For demonstration: x_{t+1} = x + dt * f(x, u) + // Adjust as needed if your model is purely discrete. + + Eigen::VectorXd x_dot = getContinuousDynamics(state, control); + return state + x_dot * timestep_; +} + +// ---------------------------------------------------------------------------- +// getStateJacobian +// ---------------------------------------------------------------------------- +Eigen::MatrixXd NeuralDynamicalSystem::getStateJacobian(const Eigen::VectorXd& state, + const Eigen::VectorXd& control) const +{ + // Placeholder approach #1: Identity, as a quick stub + // return Eigen::MatrixXd::Identity(state_dim_, state_dim_); + + // Placeholder approach #2: zero + // return Eigen::MatrixXd::Zero(state_dim_, state_dim_); + + // Real approach: use finite difference or PyTorch autograd. + // For illustration, let's do a naive finite-difference: + const double eps = 1e-6; + Eigen::MatrixXd A(state_dim_, state_dim_); + + // Baseline + Eigen::VectorXd f0 = getContinuousDynamics(state, control); + + for (int i = 0; i < state_dim_; ++i) { + Eigen::VectorXd perturbed = state; + perturbed(i) += eps; + + Eigen::VectorXd f_pert = getContinuousDynamics(perturbed, control); + A.col(i) = (f_pert - f0) / eps; + } + return A; +} + +// ---------------------------------------------------------------------------- +// getControlJacobian +// ---------------------------------------------------------------------------- +Eigen::MatrixXd NeuralDynamicalSystem::getControlJacobian(const Eigen::VectorXd& state, + const Eigen::VectorXd& control) const +{ + // Similar naive finite-difference: + const double eps = 1e-6; + Eigen::MatrixXd B(state_dim_, control_dim_); + + // Baseline + Eigen::VectorXd f0 = getContinuousDynamics(state, control); + + for (int j = 0; j < control_dim_; ++j) { + Eigen::VectorXd ctrl_pert = control; + ctrl_pert(j) += eps; + + Eigen::VectorXd f_pert = getContinuousDynamics(state, ctrl_pert); + B.col(j) = (f_pert - f0) / eps; + } + return B; +} + +// ---------------------------------------------------------------------------- +// Hessians (placeholders) +// ---------------------------------------------------------------------------- +Eigen::MatrixXd NeuralDynamicalSystem::getStateHessian(const Eigen::VectorXd& /*state*/, + const Eigen::VectorXd& /*control*/) const +{ + // Typically a 3D object. Return zero or implement if needed. + return Eigen::MatrixXd::Zero(state_dim_ * state_dim_, state_dim_); +} + +Eigen::MatrixXd NeuralDynamicalSystem::getControlHessian(const Eigen::VectorXd& /*state*/, + const Eigen::VectorXd& /*control*/) const +{ + return Eigen::MatrixXd::Zero(control_dim_ * control_dim_, state_dim_); +} + +Eigen::MatrixXd NeuralDynamicalSystem::getCrossHessian(const Eigen::VectorXd& /*state*/, + const Eigen::VectorXd& /*control*/) const +{ + return Eigen::MatrixXd::Zero(state_dim_ * control_dim_, state_dim_); +} + +// ---------------------------------------------------------------------------- +// eigenToTorch / torchToEigen +// ---------------------------------------------------------------------------- +torch::Tensor NeuralDynamicalSystem::eigenToTorch(const Eigen::VectorXd& eigen_vec, + bool requires_grad) const +{ + // Shape [1, size] + // If you prefer shape [size] without batch-dim, adjust accordingly. + auto tensor = torch::from_blob( + const_cast(eigen_vec.data()), + {1, static_cast(eigen_vec.size())}, + torch::TensorOptions().dtype(torch::kDouble) + ).clone(); // .clone() to own memory + + // Move to device + tensor = tensor.to(device_); + + // optionally set requires_grad + if (requires_grad) { + tensor.set_requires_grad(true); + } + return tensor; +} + +Eigen::VectorXd NeuralDynamicalSystem::torchToEigen(const torch::Tensor& tensor) const +{ + // Expect shape [1, state_dim_] or [1, control_dim_]. We'll read out the second dimension. + auto cpu_tensor = tensor.to(torch::kCPU).contiguous(); + if (cpu_tensor.dim() != 2) { + throw std::runtime_error("torchToEigen: expected a 2D tensor with batch-dim"); + } + auto rows = cpu_tensor.size(0); + auto cols = cpu_tensor.size(1); + + // For a single sample, rows == 1. We'll produce an Eigen vector of length = cols. + Eigen::VectorXd eigen_vec(cols); + std::memcpy(eigen_vec.data(), cpu_tensor.data_ptr(), sizeof(double)*cols); + return eigen_vec; +} + +} // namespace cddp diff --git a/src/cddp_core/torch_dynamical_system.cpp b/src/cddp_core/torch_dynamical_system.cpp deleted file mode 100644 index 9d1a894..0000000 --- a/src/cddp_core/torch_dynamical_system.cpp +++ /dev/null @@ -1,182 +0,0 @@ -/* - 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. -*/ - -#include "cddp_core/torch_dynamical_system.hpp" - -namespace cddp { - -TorchDynamicalSystem::TorchDynamicalSystem( - int state_dim, - int control_dim, - double timestep, - std::string integration_type, - std::shared_ptr model, - bool use_gpu) - : DynamicalSystem(state_dim, control_dim, timestep, integration_type), - model_(model), - use_gpu_(use_gpu && torch::cuda::is_available()), - device_(use_gpu_ ? torch::kCUDA : torch::kCPU) { - - if (use_gpu_ != use_gpu) { - std::cout << "Warning: GPU requested but not available. Using CPU instead." << std::endl; - } - - if (use_gpu_) { - model_->to(device_); - } -} - -torch::Tensor TorchDynamicalSystem::eigenToTorch( - const Eigen::VectorXd& eigen_vec, - bool requires_grad) const { - - auto tensor = torch::from_blob( - const_cast(eigen_vec.data()), - {1, eigen_vec.size()}, - torch::kFloat64 - ).clone().to(device_); // Move to correct device immediately after creation - - tensor.requires_grad_(requires_grad); - return tensor; -} - -Eigen::VectorXd TorchDynamicalSystem::torchToEigen( - const torch::Tensor& tensor) const { - - auto cpu_tensor = tensor.to(torch::kCPU).contiguous(); - Eigen::VectorXd eigen_vec(cpu_tensor.size(1)); - std::memcpy(eigen_vec.data(), - cpu_tensor.data_ptr(), - sizeof(double) * cpu_tensor.numel()); - return eigen_vec; -} - -Eigen::VectorXd TorchDynamicalSystem::getContinuousDynamics( - const Eigen::VectorXd& state, - const Eigen::VectorXd& control) const { - - torch::NoGradGuard no_grad; - - auto state_tensor = eigenToTorch(state); - auto control_tensor = eigenToTorch(control); - - auto output = model_->forward({state_tensor, control_tensor}); - return torchToEigen(output); -} - -Eigen::MatrixXd TorchDynamicalSystem::getStateJacobian( - const Eigen::VectorXd& state, - const Eigen::VectorXd& control) const { - - auto state_tensor = eigenToTorch(state, true); // requires_grad = true - auto control_tensor = eigenToTorch(control); - - auto output = model_->forward({state_tensor, control_tensor}); - - Eigen::MatrixXd jacobian(state_dim_, state_dim_); - for (int i = 0; i < output.size(1); i++) { - auto grad = torch::autograd::grad( - {output[0][i]}, - {state_tensor}, - {}, - true, - true)[0]; - - Eigen::VectorXd grad_vec = torchToEigen(grad); - jacobian.row(i) = grad_vec; - } - - return jacobian; -} - -Eigen::MatrixXd TorchDynamicalSystem::getControlJacobian( - const Eigen::VectorXd& state, - const Eigen::VectorXd& control) const { - - auto state_tensor = eigenToTorch(state); - auto control_tensor = eigenToTorch(control, true); // requires_grad = true - - auto output = model_->forward({state_tensor, control_tensor}); - - Eigen::MatrixXd jacobian(state_dim_, control_dim_); - for (int i = 0; i < output.size(1); i++) { - auto grad = torch::autograd::grad( - {output[0][i]}, - {control_tensor}, - {}, - true, - true)[0]; - - Eigen::VectorXd grad_vec = torchToEigen(grad); - jacobian.row(i) = grad_vec; - } - - return jacobian; -} - -std::vector TorchDynamicalSystem::getBatchDynamics( - const std::vector& states, - const std::vector& controls) const { - - torch::NoGradGuard no_grad; - - // Create batched tensors - auto state_batch = torch::zeros( - {static_cast(states.size()), state_dim_}, - torch::kFloat64 - ).to(device_); - - auto control_batch = torch::zeros( - {static_cast(controls.size()), control_dim_}, - torch::kFloat64 - ).to(device_); - - // Copy data to tensors - for (size_t i = 0; i < states.size(); ++i) { - state_batch[i] = eigenToTorch(states[i])[0]; - control_batch[i] = eigenToTorch(controls[i])[0]; - } - - // Forward pass - auto output_batch = model_->forward({state_batch, control_batch}); - - // Convert back to vector of Eigen vectors - std::vector results; - results.reserve(states.size()); - - for (int i = 0; i < output_batch.size(0); ++i) { - results.push_back(torchToEigen(output_batch[i].unsqueeze(0))); - } - - return results; -} - -// Optional: Implement Hessian computations if needed -Eigen::MatrixXd TorchDynamicalSystem::getStateHessian( - const Eigen::VectorXd& state, - const Eigen::VectorXd& control) const { - // TODO: Implement second-order derivatives if needed - return Eigen::MatrixXd::Zero(state_dim_ * state_dim_, state_dim_); -} - -Eigen::MatrixXd TorchDynamicalSystem::getControlHessian( - const Eigen::VectorXd& state, - const Eigen::VectorXd& control) const { - // TODO: Implement second-order derivatives if needed - return Eigen::MatrixXd::Zero(state_dim_ * control_dim_, control_dim_); -} - -} // namespace cddp \ No newline at end of file diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index cb85ec0..0d55ac8 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -94,6 +94,17 @@ gtest_discover_tests(test_logcddp_core) add_executable(test_animation test_animation.cpp) target_link_libraries(test_animation gtest gmock gtest_main cddp) +# Test for torch +if (CDDP_CPP_TORCH) + add_executable(test_torch test_torch.cpp) + target_link_libraries(test_torch gtest gmock gtest_main cddp) + gtest_discover_tests(test_torch) + + add_executable(test_neural_pendulum dynamics_model/test_neural_pendulum.cpp) + target_link_libraries(test_neural_pendulum gtest gmock gtest_main cddp) + gtest_discover_tests(test_neural_pendulum) +endif() + if (CDDP_CPP_SQP) add_executable(test_sqp sqp_core/test_sqp.cpp) target_link_libraries(test_sqp gtest gmock gtest_main cddp) @@ -121,15 +132,4 @@ if (CDDP_CPP_GUROBI) add_executable(test_qp_solvers test_qp_solvers.cpp) target_link_libraries(test_qp_solvers gtest gmock gtest_main cddp) gtest_discover_tests(test_qp_solvers) -endif() - -# Test for torch -if (CDDP_CPP_TORCH) - add_executable(test_torch test_torch.cpp) - target_link_libraries(test_torch gtest gmock gtest_main cddp) - gtest_discover_tests(test_torch) - - add_executable(test_torch_pendulum dynamics_model/test_torch_pendulum.cpp) - target_link_libraries(test_torch_pendulum gtest gmock gtest_main cddp) - gtest_discover_tests(test_torch_pendulum) endif() \ No newline at end of file diff --git a/tests/dynamics_model/test_neural_cartpole.cpp b/tests/dynamics_model/test_neural_cartpole.cpp new file mode 100644 index 0000000..e69de29 diff --git a/tests/dynamics_model/test_neural_pendulum.cpp b/tests/dynamics_model/test_neural_pendulum.cpp new file mode 100644 index 0000000..8a954e6 --- /dev/null +++ b/tests/dynamics_model/test_neural_pendulum.cpp @@ -0,0 +1,299 @@ +/* + 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. +*/ + +#include +#include +#include +#include "gmock/gmock.h" +#include "gtest/gtest.h" + +#include "cddp.hpp" +#include "cddp_core/neural_dynamical_system.hpp" + +namespace plt = matplotlibcpp; +namespace fs = std::filesystem; + +using namespace cddp; + +class PendulumModel : public DynamicsModelInterface { +public: + PendulumModel(double length = 1.0, double mass = 1.0, double damping = 0.0); + + torch::Tensor forward(std::vector inputs) override; + +private: + void initialize_weights(); + + double length_, mass_, damping_; + torch::nn::Linear linear1{nullptr}, linear2{nullptr}, linear3{nullptr}; + torch::Device device_; +}; + +PendulumModel::PendulumModel(double length, double mass, double damping) + : length_(length), mass_(mass), damping_(damping), + device_(torch::cuda::is_available() ? torch::kCUDA : torch::kCPU) +{ + // Create linear layers + linear1 = register_module("linear1", torch::nn::Linear(3, 32)); + linear2 = register_module("linear2", torch::nn::Linear(32, 32)); + linear3 = register_module("linear3", torch::nn::Linear(32, 2)); + + // Move to device and set dtype + this->to(device_); + if (device_.is_cpu()) { + this->to(torch::kFloat64); + } else { + this->to(torch::kFloat32); + } + initialize_weights(); +} + +void PendulumModel::initialize_weights() +{ + torch::NoGradGuard no_grad; + + double angle_scale = 0.1; + double velocity_scale = 0.1; + + // Example: manually set a few weights in linear3 + auto w = linear3->weight.data(); + w[0][0] = angle_scale; + w[1][1] = velocity_scale; + linear3->weight.data() = w; + + auto b = linear3->bias.data(); + b[0] = 0.0; + b[1] = -9.81 / length_ * angle_scale; // a quick guess + linear3->bias.data() = b; +} + +torch::Tensor PendulumModel::forward(std::vector inputs) +{ + // Expect 2 inputs: [state, control] + auto state = inputs[0].to(device_); + auto control = inputs[1].to(device_); + + if (device_.is_cuda()) { + // If on GPU, we use float32 + state = state.to(torch::kFloat32); + control = control.to(torch::kFloat32); + } + + // Concatenate along dim=1 if shapes are [batch_size, 2] and [batch_size, 1] + auto x = torch::cat({state, control}, /*dim=*/1); + x = torch::tanh(linear1(x)); + x = torch::tanh(linear2(x)); + x = linear3(x); + + if (device_.is_cuda()) { + x = x.to(torch::kFloat64); // move back to double if desired + } + return x; +} + +void printTensorInfo(const torch::Tensor& tensor, const std::string& name) +{ + std::cout << name << ":\n" + << " - shape: [" << tensor.sizes() << "]\n" + << " - dtype: " << tensor.dtype() << "\n" + << " - device: " << tensor.device() << "\n" + << " - values: " << tensor << "\n" << std::endl; +} + +void printVectorInfo(const Eigen::VectorXd& vec, const std::string& name) +{ + std::cout << name << ":\n" + << " - size: " << vec.size() << "\n" + << " - values: " << vec.transpose() << "\n" << std::endl; +} + +//-----------------------------------------// +// TEST SUITE: TorchPendulumTest +//-----------------------------------------// +TEST(TorchPendulumTest, DiscreteDynamics) +{ + // Parameters + double timestep = 0.01; + double length = 1.0; + double mass = 1.0; + double damping = 0.0; + std::string integration_type = "rk4"; + + cddp::Pendulum analytical_pendulum(timestep, length, mass, damping, integration_type); + + auto model = std::make_shared(length, mass, damping); + torch::Device device(torch::cuda::is_available() ? torch::kCUDA : torch::kCPU); + NeuralDynamicalSystem torch_pendulum(/*state_dim=*/2, /*control_dim=*/1, + timestep, integration_type, model, device); + + // Create a test state/control + Eigen::VectorXd test_state(2); + test_state << M_PI / 4, 0.0; // 45 deg, no velocity + Eigen::VectorXd test_control(1); + test_control << 0.0; + + // Analytical next state + auto analytical_next = analytical_pendulum.getDiscreteDynamics(test_state, test_control); + printVectorInfo(analytical_next, "Analytical Next State"); + + // Torch next state + auto torch_next = torch_pendulum.getDiscreteDynamics(test_state, test_control); + printVectorInfo(torch_next, "Torch Next State"); + + // Compare errors + double error = (analytical_next - torch_next).norm(); + std::cout << "L2 error: " << error << std::endl; + + // Basic tests + ASSERT_EQ(torch_pendulum.getStateDim(), 2); + ASSERT_EQ(torch_pendulum.getControlDim(), 1); + ASSERT_DOUBLE_EQ(torch_pendulum.getTimestep(), timestep); + ASSERT_EQ(torch_pendulum.getIntegrationType(), integration_type); + + // Check if the error is within a tolerance + // (This is arbitrary; you may need a looser or tighter tolerance) + EXPECT_NEAR(error, 0.0, 0.1) + << "Discrete dynamics: model deviance too large"; + + // Test Jacobians + auto A = torch_pendulum.getStateJacobian(test_state, test_control); + auto B = torch_pendulum.getControlJacobian(test_state, test_control); + + std::cout << "State Jacobian A:\n" << A << std::endl; + std::cout << "Control Jacobian B:\n" << B << std::endl; + + // Verify shapes + ASSERT_EQ(A.rows(), 2); + ASSERT_EQ(A.cols(), 2); + ASSERT_EQ(B.rows(), 2); + ASSERT_EQ(B.cols(), 1); +} + +/** + * @brief Demonstrates a small training loop that tries to fit the Torch model + * to data from an analytical pendulum's discrete dynamics. + */ +TEST(NeuralPendulumTest, Training) +{ + // Parameters + double timestep = 0.01; + double length = 1.0; + double mass = 1.0; + double damping = 0.1; + + // Create an analytical pendulum and a Torch model + cddp::Pendulum analytical_pendulum(timestep, length, mass, damping, "rk4"); + auto model = std::make_shared(length, mass, damping); + bool use_gpu = torch::cuda::is_available(); + + torch::optim::Adam optimizer(model->parameters(), torch::optim::AdamOptions(1e-3)); + int num_epochs = 100; + int batch_size = 32; + + // Generate training data + int num_samples = 1000; + auto cpu_options = torch::TensorOptions().dtype(torch::kFloat64).device(torch::kCPU); + auto state_tensor = torch::zeros({num_samples, 2}, cpu_options); + auto control_tensor = torch::zeros({num_samples, 1}, cpu_options); + auto next_state_tensor = torch::zeros({num_samples, 2}, cpu_options); + + for (int i = 0; i < num_samples; ++i) { + Eigen::VectorXd state(2); + state << (2.0 * rand() / RAND_MAX - 1.0) * M_PI, + (2.0 * rand() / RAND_MAX - 1.0) * 5.0; + Eigen::VectorXd control(1); + control << (2.0 * rand() / RAND_MAX - 1.0) * 2.0; + + auto next_state = analytical_pendulum.getDiscreteDynamics(state, control); + + // Copy to torch Tensors + state_tensor[i] = torch::from_blob(state.data(), {2}, cpu_options).clone(); + control_tensor[i] = torch::from_blob(control.data(), {1}, cpu_options).clone(); + next_state_tensor[i] = torch::from_blob(next_state.data(), {2}, cpu_options).clone(); + } + + // Move to GPU if available + torch::Device device(use_gpu ? torch::kCUDA : torch::kCPU); + if (use_gpu) { + state_tensor = state_tensor.to(device).to(torch::kFloat32); + control_tensor = control_tensor.to(device).to(torch::kFloat32); + next_state_tensor = next_state_tensor.to(device).to(torch::kFloat32); + } else { + state_tensor = state_tensor.to(device); + control_tensor = control_tensor.to(device); + next_state_tensor = next_state_tensor.to(device); + } + + // Training loop + for (int epoch = 0; epoch < num_epochs; ++epoch) { + double total_loss = 0.0; + int num_batches = num_samples / batch_size; + + for (int batch = 0; batch < num_batches; ++batch) { + auto indices = torch::randperm(num_samples, + torch::TensorOptions().dtype(torch::kLong).device(device)); + indices = indices.slice(0, 0, batch_size); + + auto batch_states = state_tensor.index_select(0, indices); + auto batch_controls = control_tensor.index_select(0, indices); + auto batch_next_states = next_state_tensor.index_select(0, indices); + + optimizer.zero_grad(); + auto pred_next_states = model->forward({batch_states, batch_controls}); + if (use_gpu) { + pred_next_states = pred_next_states.to(torch::kFloat32); + } + auto loss = torch::mse_loss(pred_next_states, batch_next_states); + loss.backward(); + optimizer.step(); + + total_loss += loss.item(); + } + + if (epoch % 10 == 0) { + std::cout << "Epoch " << epoch + << ", Avg Loss: " << total_loss / num_batches << std::endl; + } + } + + // Test final performance + Eigen::VectorXd test_state(2); + test_state << M_PI/4, 0.0; + Eigen::VectorXd test_control(1); + test_control << 0.0; + + // Move test inputs to torch + auto cpu_options_64 = torch::TensorOptions().dtype(torch::kFloat64).device(torch::kCPU); + auto test_state_tensor = torch::from_blob(test_state.data(), {1, 2}, cpu_options_64).clone(); + auto test_control_tensor = torch::from_blob(test_control.data(), {1, 1}, cpu_options_64).clone(); + if (use_gpu) { + test_state_tensor = test_state_tensor.to(device).to(torch::kFloat32); + test_control_tensor = test_control_tensor.to(device).to(torch::kFloat32); + } + auto pred = model->forward({test_state_tensor, test_control_tensor}); + pred = pred.to(torch::kCPU).to(torch::kFloat64); // back to CPU/double + + // Copy to Eigen + Eigen::VectorXd pred_eigen(2); + std::memcpy(pred_eigen.data(), pred.data_ptr(), sizeof(double) * 2); + + // Compare to analytical + auto analytical_next = analytical_pendulum.getDiscreteDynamics(test_state, test_control); + double error = (analytical_next - pred_eigen).norm(); + std::cout << "Test error: " << error << std::endl; + EXPECT_LT(error, 0.1); +} + diff --git a/tests/dynamics_model/test_neural_quadrotor.cpp b/tests/dynamics_model/test_neural_quadrotor.cpp new file mode 100644 index 0000000..e69de29 diff --git a/tests/dynamics_model/test_torch_pendulum.cpp b/tests/dynamics_model/test_torch_pendulum.cpp deleted file mode 100644 index 57a39d6..0000000 --- a/tests/dynamics_model/test_torch_pendulum.cpp +++ /dev/null @@ -1,346 +0,0 @@ -/* - 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. -*/ - -#include -#include -#include - -#include "gmock/gmock.h" -#include "gtest/gtest.h" -#include "cddp.hpp" -#include "cddp_core/torch_dynamical_system.hpp" - -namespace plt = matplotlibcpp; -namespace fs = std::filesystem; -using namespace cddp; - -class PendulumModel : public DynamicsModelInterface { -public: - PendulumModel(double length = 1.0, double mass = 1.0, double damping = 0.0) - : length_(length), mass_(mass), damping_(damping), - device_(torch::cuda::is_available() ? torch::kCUDA : torch::kCPU) { - - // Create linear layers - linear1 = register_module("linear1", torch::nn::Linear(3, 32)); - linear2 = register_module("linear2", torch::nn::Linear(32, 32)); - linear3 = register_module("linear3", torch::nn::Linear(32, 2)); - - // Move to appropriate device and dtype - this->to(device_); - if (device_.is_cpu()) { - this->to(torch::kFloat64); - } else { - this->to(torch::kFloat32); // Use Float32 for GPU - } - initialize_weights(); - } - - torch::Tensor forward(std::vector inputs) override { - auto state = inputs[0].to(device_); - auto control = inputs[1].to(device_); - - if (device_.is_cuda()) { - state = state.to(torch::kFloat32); - control = control.to(torch::kFloat32); - } - - auto x = torch::cat({state, control}, /*dim=*/1); - x = torch::tanh(linear1(x)); - x = torch::tanh(linear2(x)); - x = linear3(x); - - if (device_.is_cuda()) { - x = x.to(torch::kFloat64); - } - - return x; - } - -private: - void initialize_weights() { - torch::NoGradGuard no_grad; - - double angle_scale = 0.1; - double velocity_scale = 0.1; - - auto w = linear3->weight.data(); - w[0][0] = angle_scale; - w[1][1] = velocity_scale; - linear3->weight.data() = w.to(torch::kFloat64); - - auto b = linear3->bias.data(); - b[0] = 0.0; - b[1] = -9.81 / length_ * angle_scale; - linear3->bias.data() = b.to(torch::kFloat64); - } - - torch::Tensor createTensor(const Eigen::VectorXd& eigen_vec) { - // First create tensor on CPU - auto cpu_tensor = torch::from_blob( - const_cast(eigen_vec.data()), - {eigen_vec.size()}, - torch::kFloat64 - ).clone(); // Clone to own the memory - - // Then move to target device - return cpu_tensor.to(device_); - } - - double length_, mass_, damping_; - torch::nn::Linear linear1{nullptr}, linear2{nullptr}, linear3{nullptr}; - torch::Device device_; -}; - -void printTensorInfo(const torch::Tensor& tensor, const std::string& name) { - std::cout << name << ":\n" - << " - shape: [" << tensor.sizes() << "]\n" - << " - dtype: " << tensor.dtype() << "\n" - << " - device: " << tensor.device() << "\n" - << " - values: " << tensor << "\n" << std::endl; -} - -void printVectorInfo(const Eigen::VectorXd& vec, const std::string& name) { - std::cout << name << ":\n" - << " - size: " << vec.size() << "\n" - << " - values: " << vec.transpose() << "\n" << std::endl; -} - -// TEST(TorchPendulumTest, DiscreteDynamics) { -// // Parameters -// double timestep = 0.01; -// double length = 1.0; -// double mass = 1.0; -// double damping = 0.0; -// std::string integration_type = "rk4"; - -// // Create analytical pendulum for comparison -// cddp::Pendulum analytical_pendulum(timestep, length, mass, damping, integration_type); - -// // Create torch model and system -// auto model = std::make_shared(length, mass, damping); -// bool use_gpu = torch::cuda::is_available(); -// TorchDynamicalSystem torch_pendulum(2, 1, timestep, integration_type, model, use_gpu); - -// // Create device object -// torch::Device device(use_gpu ? torch::kCUDA : torch::kCPU); - -// // Initial test state and control -// Eigen::VectorXd test_state(2); -// test_state << M_PI/4, 0.0; // 45 degrees, no initial velocity -// Eigen::VectorXd test_control(1); -// test_control << 0.0; // No control input - -// // Print initial conditions -// std::cout << "\nInitial Test Conditions:" << std::endl; -// printVectorInfo(test_state, "Test State"); -// printVectorInfo(test_control, "Test Control"); - -// // Get analytical prediction -// auto analytical_next = analytical_pendulum.getDiscreteDynamics(test_state, test_control); -// std::cout << "\nAnalytical Model Prediction:" << std::endl; -// printVectorInfo(analytical_next, "Analytical Next State"); - -// // Get torch prediction -// auto torch_next = torch_pendulum.getDiscreteDynamics(test_state, test_control); -// std::cout << "\nTorch Model Prediction:" << std::endl; -// printVectorInfo(torch_next, "Torch Next State"); - -// // Print error analysis -// double error = (analytical_next - torch_next).norm(); -// std::cout << "\nError Analysis:" << std::endl; -// std::cout << "Total Error (L2 norm): " << error << std::endl; -// std::cout << "Component-wise errors:" << std::endl; -// std::cout << "Position error: " << std::abs(analytical_next[0] - torch_next[0]) << std::endl; -// std::cout << "Velocity error: " << std::abs(analytical_next[1] - torch_next[1]) << std::endl; - -// // Basic assertions -// ASSERT_EQ(torch_pendulum.getStateDim(), 2); -// ASSERT_EQ(torch_pendulum.getControlDim(), 1); -// ASSERT_DOUBLE_EQ(torch_pendulum.getTimestep(), 0.01); -// ASSERT_EQ(torch_pendulum.getIntegrationType(), "rk4"); - -// // Prediction accuracy tests -// EXPECT_NEAR((analytical_next - torch_next).norm(), 0.0, 0.1) -// << "Total prediction error too large"; -// EXPECT_NEAR(analytical_next[0], torch_next[0], 0.1) -// << "Position prediction error too large"; -// EXPECT_NEAR(analytical_next[1], torch_next[1], 0.1) -// << "Velocity prediction error too large"; - -// // Test with different initial conditions -// std::vector test_states = { -// (Eigen::VectorXd(2) << 0.0, 0.0).finished(), // At rest -// (Eigen::VectorXd(2) << M_PI/2, 0.0).finished(), // 90 degrees -// (Eigen::VectorXd(2) << 0.0, 1.0).finished(), // With initial velocity -// (Eigen::VectorXd(2) << M_PI/4, -1.0).finished() // 45 degrees with negative velocity -// }; - -// std::cout << "\nTesting multiple initial conditions:" << std::endl; -// for (const auto& state : test_states) { -// std::cout << "\nTesting state: " << state.transpose() << std::endl; - -// auto analytical = analytical_pendulum.getDiscreteDynamics(state, test_control); -// auto torch = torch_pendulum.getDiscreteDynamics(state, test_control); - -// double test_error = (analytical - torch).norm(); -// std::cout << "Error: " << test_error << std::endl; - -// EXPECT_NEAR(test_error, 0.0, 0.1) -// << "Large error for initial state: " << state.transpose(); -// } - -// // Test Jacobian computation -// auto state_jacobian = torch_pendulum.getStateJacobian(test_state, test_control); -// auto control_jacobian = torch_pendulum.getControlJacobian(test_state, test_control); - -// std::cout << "\nJacobian Analysis:" << std::endl; -// std::cout << "State Jacobian:\n" << state_jacobian << std::endl; -// std::cout << "Control Jacobian:\n" << control_jacobian << std::endl; - -// // Verify Jacobian dimensions -// ASSERT_EQ(state_jacobian.rows(), 2); -// ASSERT_EQ(state_jacobian.cols(), 2); -// ASSERT_EQ(control_jacobian.rows(), 2); -// ASSERT_EQ(control_jacobian.cols(), 1); -// } - - -TEST(TorchPendulumTest, Training) { - // Parameters - double timestep = 0.01; - double length = 1.0; - double mass = 1.0; - double damping = 0.1; - - // Create pendulums - cddp::Pendulum analytical_pendulum(timestep, length, mass, damping, "rk4"); - auto model = std::make_shared(length, mass, damping); - bool use_gpu = torch::cuda::is_available(); - - torch::optim::Adam optimizer(model->parameters(), torch::optim::AdamOptions(1e-3)); - int num_epochs = 100; - int batch_size = 32; - - // Generate training data - int num_samples = 1000; - std::vector states, next_states; - std::vector controls; - - // Create tensors initially on CPU with Float64 - auto cpu_options = torch::TensorOptions() - .dtype(torch::kFloat64) - .device(torch::kCPU); - - auto state_tensor = torch::zeros({num_samples, 2}, cpu_options); - auto control_tensor = torch::zeros({num_samples, 1}, cpu_options); - auto next_state_tensor = torch::zeros({num_samples, 2}, cpu_options); - - // Fill tensors - for (int i = 0; i < num_samples; ++i) { - Eigen::VectorXd state(2); - state << (2.0 * (double)rand() / RAND_MAX - 1.0) * M_PI, - (2.0 * (double)rand() / RAND_MAX - 1.0) * 5.0; - - Eigen::VectorXd control(1); - control << (2.0 * (double)rand() / RAND_MAX - 1.0) * 2.0; - - states.push_back(state); - controls.push_back(control); - next_states.push_back(analytical_pendulum.getDiscreteDynamics(state, control)); - - state_tensor[i] = torch::from_blob(state.data(), {2}, cpu_options).clone(); - control_tensor[i] = torch::from_blob(control.data(), {1}, cpu_options).clone(); - next_state_tensor[i] = torch::from_blob(next_states[i].data(), {2}, cpu_options).clone(); - } - - // Move to GPU if available and convert to Float32 for GPU operations - torch::Device device(use_gpu ? torch::kCUDA : torch::kCPU); - if (use_gpu) { - state_tensor = state_tensor.to(device).to(torch::kFloat32); - control_tensor = control_tensor.to(device).to(torch::kFloat32); - next_state_tensor = next_state_tensor.to(device).to(torch::kFloat32); - } else { - state_tensor = state_tensor.to(device); - control_tensor = control_tensor.to(device); - next_state_tensor = next_state_tensor.to(device); - } - - // Training loop - for (int epoch = 0; epoch < num_epochs; ++epoch) { - double total_loss = 0.0; - int num_batches = num_samples / batch_size; - - for (int batch = 0; batch < num_batches; ++batch) { - // Create indices with appropriate dtype for GPU - auto indices = torch::randperm(num_samples, - torch::TensorOptions() - .dtype(torch::kLong) // Important: indices must be Long - .device(device)); - indices = indices.slice(0, 0, batch_size); - - auto batch_states = state_tensor.index_select(0, indices); - auto batch_controls = control_tensor.index_select(0, indices); - auto batch_next_states = next_state_tensor.index_select(0, indices); - - optimizer.zero_grad(); - auto pred_next_states = model->forward({batch_states, batch_controls}); - - if (use_gpu) { - pred_next_states = pred_next_states.to(torch::kFloat32); - } - - auto loss = torch::mse_loss(pred_next_states, batch_next_states); - loss.backward(); - optimizer.step(); - - total_loss += loss.item(); - } - - if (epoch % 10 == 0) { - std::cout << "Epoch " << epoch - << ", Average Loss: " << total_loss / num_batches - << std::endl; - } - } - - // Testing - Eigen::VectorXd test_state(2); - test_state << M_PI/4, 0.0; - Eigen::VectorXd test_control(1); - test_control << 0.0; - - // Create test tensors on CPU first - auto test_state_tensor = torch::from_blob(test_state.data(), {1, 2}, cpu_options).clone(); - auto test_control_tensor = torch::from_blob(test_control.data(), {1, 1}, cpu_options).clone(); - - if (use_gpu) { - test_state_tensor = test_state_tensor.to(device).to(torch::kFloat32); - test_control_tensor = test_control_tensor.to(device).to(torch::kFloat32); - } else { - test_state_tensor = test_state_tensor.to(device); - test_control_tensor = test_control_tensor.to(device); - } - - auto pred = model->forward({test_state_tensor, test_control_tensor}); - pred = pred.to(torch::kCPU).to(torch::kFloat64); // Move back to CPU and double precision - - Eigen::VectorXd pred_eigen(2); - std::memcpy(pred_eigen.data(), pred.data_ptr(), sizeof(double) * 2); - - auto analytical_next = analytical_pendulum.getDiscreteDynamics(test_state, test_control); - double error = (analytical_next - pred_eigen).norm(); - EXPECT_LT(error, 0.1); -} \ No newline at end of file