TestLibtorchRLCore.C

Go to the documentation of this file.

//* This file is part of the MOOSE framework
//* https://mooseframework.inl.gov
//*
//* All rights reserved, see COPYRIGHT for full restrictions
//* https://github.com/idaholab/moose/blob/master/COPYRIGHT
//*
//* Licensed under LGPL 2.1, please see LICENSE for details
//* https://www.gnu.org/licenses/lgpl-2.1.html
 
#ifdef MOOSE_LIBTORCH_ENABLED
 
#include "gtest/gtest.h"
 
#include "LibtorchActorNeuralNet.h"
#include "LibtorchArtificialNeuralNet.h"
#include "LibtorchObservationHistoryHelper.h"
#include "LibtorchRandomUtils.h"
#include "LibtorchRLMiniBatchSampler.h"
#include "LibtorchRLPPOLoss.h"
#include "LibtorchRLTrajectoryAssembler.h"
#include "LibtorchRLTrajectoryBuffer.h"
#include "LibtorchRLValueEstimator.h"
 
#include <cmath>
 
namespace
{
 
std::vector<Real>
linearSequence(const unsigned int size, const Real offset, const Real slope)
{
  std::vector<Real> values(size);
  for (const auto i : make_range(size))
    values[i] = offset + slope * i;
  return values;
}
 
LibtorchRLTrajectoryBuffer::Trajectory
trajectoryWithRewards(const std::vector<Real> & rewards)
{
  LibtorchRLTrajectoryBuffer::Trajectory trajectory;
  trajectory.observations.assign(rewards.size(), {0.0});
  trajectory.next_observations.assign(rewards.size(), {0.0});
  trajectory.actions.assign(rewards.size(), {0.0});
  trajectory.log_probabilities.assign(rewards.size(), {0.0});
  trajectory.rewards = rewards;
  return trajectory;
}
 
TEST(LibtorchRLCoreTest, ObservationHistoryStacksCurrentAndTrajectoryData)
{
  LibtorchObservationHistoryHelper history(3);
 
  const auto expanded_observation_factors = history.expandObservationFactors({0.5, 2.0});
  EXPECT_EQ(expanded_observation_factors, std::vector<Real>({0.5, 2.0, 0.5, 2.0, 0.5, 2.0}));
 
  std::vector<std::vector<Real>> old_observations;
  history.initializeHistory({1.0, 6.0}, old_observations);
 
  const auto stacked_current = history.stackCurrentObservation({3.0, 2.0}, old_observations);
  EXPECT_EQ(stacked_current, std::vector<Real>({3.0, 2.0, 1.0, 6.0, 1.0, 6.0}));
 
  std::vector<std::vector<Real>> component_trajectories = {{1.0, 2.0, 3.0}, {4.0, 5.0, 6.0}};
  const auto stacked_trajectory = history.stackTrajectoryObservation(component_trajectories, 2);
  EXPECT_EQ(stacked_trajectory, std::vector<Real>({3.0, 6.0, 2.0, 5.0, 1.0, 4.0}));
}
 
TEST(LibtorchRLCoreTest, ObservationHistoryRejectsMalformedStoredHistory)
{
  LibtorchObservationHistoryHelper history(3);
 
  EXPECT_THROW(
      {
        try
        {
          history.stackCurrentObservation({3.0, 2.0}, {{1.0, 6.0}});
        }
        catch (const std::exception & e)
        {
          EXPECT_EQ(std::string(e.what()),
                    "Observation history must contain 2 stored entries, but got 1.");
          throw;
        }
      },
      std::exception);
}
 
TEST(LibtorchRLCoreTest, ObservationHistoryRejectsZeroTimesteps)
{
  EXPECT_THROW(
      {
        try
        {
          LibtorchObservationHistoryHelper history(0);
          static_cast<void>(history);
        }
        catch (const std::exception & e)
        {
          EXPECT_EQ(std::string(e.what()),
                    "Observation history requires at least one input timestep.");
          throw;
        }
      },
      std::exception);
}
 
TEST(LibtorchRLCoreTest, ValueEstimatorComputesGAETargets)
{
  Moose::LibtorchArtificialNeuralNet value_network("value", 1, 1, {}, {"linear"});
  auto value_params = value_network.named_parameters();
  value_params[0].value().data().fill_(1.0);
  value_params[1].value().data().fill_(0.0);
 
  LibtorchRLTrajectoryBuffer::Trajectory trajectory;
  trajectory.observations = {{1.0}, {2.0}};
  trajectory.next_observations = {{2.0}, {3.0}};
  trajectory.rewards = {0.5, 1.0};
 
  LibtorchRLValueEstimator estimator(0.9, 0.95);
  const auto targets = estimator.estimate(trajectory, value_network);
 
  ASSERT_EQ(targets.advantages.size(), 2u);
  ASSERT_EQ(targets.value_targets.size(), 2u);
  EXPECT_NEAR(targets.advantages[0], 2.7535, 1e-12);
  EXPECT_NEAR(targets.advantages[1], 1.7, 1e-12);
  EXPECT_NEAR(targets.value_targets[0], 3.7535, 1e-12);
  EXPECT_NEAR(targets.value_targets[1], 3.7, 1e-12);
}
 
TEST(LibtorchRLCoreTest, ValueEstimatorMasksTerminalTransitions)
{
  Moose::LibtorchArtificialNeuralNet value_network("value", 1, 1, {}, {"linear"});
  auto value_params = value_network.named_parameters();
  value_params[0].value().data().fill_(1.0);
  value_params[1].value().data().fill_(0.0);
 
  LibtorchRLTrajectoryBuffer::Trajectory trajectory;
  trajectory.observations = {{1.0}, {2.0}};
  trajectory.next_observations = {{2.0}, {3.0}};
  trajectory.rewards = {0.5, 1.0};
  trajectory.terminals = {false, true};
 
  LibtorchRLValueEstimator estimator(0.9, 0.95);
  const auto targets = estimator.estimate(trajectory, value_network);
 
  ASSERT_EQ(targets.advantages.size(), 2u);
  ASSERT_EQ(targets.value_targets.size(), 2u);
  EXPECT_NEAR(targets.advantages[0], 0.445, 1e-12);
  EXPECT_NEAR(targets.advantages[1], -1.0, 1e-12);
  EXPECT_NEAR(targets.value_targets[0], 1.445, 1e-12);
  EXPECT_NEAR(targets.value_targets[1], 1.0, 1e-12);
}
 
TEST(LibtorchRLCoreTest, PPOLossUsesStoredLogProbabilityAndValueTarget)
{
  Moose::LibtorchActorNeuralNet policy_network("policy", 1, 1, {}, {"linear"});
  policy_network.gaussianActionDistribution().meanModule()->weight.data().fill_(0.0);
  policy_network.gaussianActionDistribution().meanModule()->bias.data().fill_(0.0);
  policy_network.gaussianActionDistribution().logStd().data().fill_(0.0);
 
  Moose::LibtorchArtificialNeuralNet value_network("value", 1, 1, {}, {"linear"});
  auto value_params = value_network.named_parameters();
  value_params[0].value().data().fill_(0.0);
  value_params[1].value().data().fill_(1.0);
 
  LibtorchRLMiniBatch batch;
  batch.observations = torch::zeros({1, 1}, torch::TensorOptions().dtype(torch::kDouble));
  batch.actions = torch::zeros({1, 1}, torch::TensorOptions().dtype(torch::kDouble));
  batch.old_log_probabilities =
      torch::tensor({{-0.5 * std::log(2.0 * libMesh::pi)}},
                    torch::TensorOptions().dtype(torch::kDouble));
  batch.value_targets = torch::tensor({{1.5}}, torch::TensorOptions().dtype(torch::kDouble));
  batch.advantages = torch::tensor({{2.0}}, torch::TensorOptions().dtype(torch::kDouble));
 
  LibtorchRLPPOLoss loss(0.2, 0.01);
  const auto loss_values = loss.compute(policy_network, value_network, batch);
 
  const Real expected_entropy = 0.5 * std::log(2.0 * libMesh::pi) + 0.5;
  const Real expected_actor_loss = -(2.0 + 0.01 * expected_entropy);
  const Real expected_critic_loss = 0.25;
 
  EXPECT_NEAR(loss_values.entropy.item<Real>(), expected_entropy, 1e-12);
  EXPECT_NEAR(loss_values.actor_loss.item<Real>(), expected_actor_loss, 1e-12);
  EXPECT_NEAR(loss_values.critic_loss.item<Real>(), expected_critic_loss, 1e-12);
}
 
TEST(LibtorchRLCoreTest, MiniBatchSamplerStandardizesAdvantagesPerBatch)
{
  LibtorchRLTrajectoryBuffer::TensorBatch batch;
  batch.observations =
      torch::tensor({{0.0}, {1.0}, {2.0}, {3.0}}, torch::TensorOptions().dtype(torch::kDouble));
  batch.actions =
      torch::tensor({{0.1}, {0.2}, {0.3}, {0.4}}, torch::TensorOptions().dtype(torch::kDouble));
  batch.log_probabilities =
      torch::tensor({{-1.0}, {-1.1}, {-1.2}, {-1.3}}, torch::TensorOptions().dtype(torch::kDouble));
  batch.value_targets =
      torch::tensor({{1.0}, {2.0}, {3.0}, {4.0}}, torch::TensorOptions().dtype(torch::kDouble));
  batch.advantages =
      torch::tensor({{1.0}, {2.0}, {3.0}, {4.0}}, torch::TensorOptions().dtype(torch::kDouble));
 
  LibtorchRLMiniBatchSampler sampler;
  const auto mini_batches = sampler.sample(batch, 2, true);
 
  ASSERT_EQ(mini_batches.size(), 2u);
  for (const auto & mini_batch : mini_batches)
  {
    ASSERT_EQ(mini_batch.size(), 2);
    EXPECT_NEAR(mini_batch.advantages.mean().item<Real>(), 0.0, 1e-12);
  }
}
 
TEST(LibtorchRLCoreTest, MiniBatchSamplerUsesExplicitGeneratorForDeterministicShuffling)
{
  LibtorchRLTrajectoryBuffer::TensorBatch batch;
  batch.observations =
      torch::tensor({{0.0}, {1.0}, {2.0}, {3.0}}, torch::TensorOptions().dtype(torch::kDouble));
  batch.actions =
      torch::tensor({{0.1}, {0.2}, {0.3}, {0.4}}, torch::TensorOptions().dtype(torch::kDouble));
  batch.log_probabilities =
      torch::tensor({{-1.0}, {-1.1}, {-1.2}, {-1.3}}, torch::TensorOptions().dtype(torch::kDouble));
  batch.value_targets =
      torch::tensor({{1.0}, {2.0}, {3.0}, {4.0}}, torch::TensorOptions().dtype(torch::kDouble));
  batch.advantages =
      torch::tensor({{1.0}, {2.0}, {3.0}, {4.0}}, torch::TensorOptions().dtype(torch::kDouble));
 
  LibtorchRLMiniBatchSampler sampler;
  const auto first_batches = sampler.sample(batch, 2, false, Moose::makeLibtorchCPUGenerator(9876));
  const auto second_batches =
      sampler.sample(batch, 2, false, Moose::makeLibtorchCPUGenerator(9876));
 
  ASSERT_EQ(first_batches.size(), second_batches.size());
  for (const auto i : index_range(first_batches))
  {
    EXPECT_TRUE(torch::allclose(first_batches[i].observations,
                                second_batches[i].observations,
                                /* rtol = */ 0.0,
                                /* atol = */ 0.0));
    EXPECT_TRUE(torch::allclose(first_batches[i].actions,
                                second_batches[i].actions,
                                /* rtol = */ 0.0,
                                /* atol = */ 0.0));
  }
}
 
TEST(LibtorchRLCoreTest, TrajectoryAssemblerDownsamplesWithWindowStride)
{
  // Raw index i holds 0 + i for observations, 100 + i for actions, -i for log probabilities, and
  // 10 * i for rewards, so each assembled value identifies the raw index it came from.
  const std::vector<std::vector<Real>> observations = {linearSequence(16, 0.0, 1.0)};
  const std::vector<std::vector<Real>> actions = {linearSequence(16, 100.0, 1.0)};
  const std::vector<std::vector<Real>> log_probabilities = {linearSequence(16, 0.0, -1.0)};
  const std::vector<std::vector<Real>> rewards = {linearSequence(16, 0.0, 10.0)};
 
  LibtorchRLTrajectoryAssembler assembler(2, 5, true, false);
  const auto trajectories =
      assembler.assemble({&observations}, {&actions}, {&log_probabilities}, rewards);
 
  ASSERT_EQ(trajectories.size(), 1u);
  const auto & trajectory = trajectories[0];
  // History lags span one window, and lags before the initial entry repeat raw index 0.
  EXPECT_EQ(trajectory.observations,
            std::vector<std::vector<Real>>({{0.0, 0.0}, {5.0, 0.0}, {10.0, 5.0}}));
  EXPECT_EQ(trajectory.next_observations,
            std::vector<std::vector<Real>>({{5.0, 0.0}, {10.0, 5.0}, {15.0, 10.0}}));
  EXPECT_EQ(trajectory.actions, std::vector<std::vector<Real>>({{101.0}, {106.0}, {111.0}}));
  EXPECT_EQ(trajectory.log_probabilities,
            std::vector<std::vector<Real>>({{-1.0}, {-6.0}, {-11.0}}));
  EXPECT_EQ(trajectory.rewards, std::vector<Real>({50.0, 100.0, 150.0}));
  EXPECT_EQ(trajectory.terminals, std::vector<bool>({false, false, true}));
}
 
TEST(LibtorchRLCoreTest, TrajectoryAssemblerAveragesRewardsWithoutShift)
{
  const std::vector<std::vector<Real>> observations = {linearSequence(16, 0.0, 1.0)};
  const std::vector<std::vector<Real>> actions = {linearSequence(16, 100.0, 1.0)};
  const std::vector<std::vector<Real>> log_probabilities = {linearSequence(16, 0.0, -1.0)};
  const std::vector<std::vector<Real>> rewards = {linearSequence(16, 0.0, 10.0)};
 
  LibtorchRLTrajectoryAssembler assembler(1, 5, false, true);
  const auto trajectories =
      assembler.assemble({&observations}, {&actions}, {&log_probabilities}, rewards);
 
  ASSERT_EQ(trajectories.size(), 1u);
  const auto & trajectory = trajectories[0];
  EXPECT_EQ(trajectory.observations, std::vector<std::vector<Real>>({{0.0}, {5.0}, {10.0}}));
  EXPECT_EQ(trajectory.actions, std::vector<std::vector<Real>>({{100.0}, {105.0}, {110.0}}));
  EXPECT_EQ(trajectory.log_probabilities, std::vector<std::vector<Real>>({{0.0}, {-5.0}, {-10.0}}));
  // Means of raw indices 1-5, 6-10, and 11-15.
  EXPECT_EQ(trajectory.rewards, std::vector<Real>({30.0, 80.0, 130.0}));
}
 
TEST(LibtorchRLCoreTest, TrajectoryAssemblerCountsCompleteWindows)
{
  LibtorchRLTrajectoryAssembler assembler(1, 5, true, false);
  EXPECT_EQ(assembler.numTransitions(0), 0u);
  EXPECT_EQ(assembler.numTransitions(1), 0u);
  EXPECT_EQ(assembler.numTransitions(5), 0u);
  EXPECT_EQ(assembler.numTransitions(6), 1u);
  EXPECT_EQ(assembler.numTransitions(16), 3u);
 
  // A sample without a complete window produces no trajectory.
  const std::vector<std::vector<Real>> observations = {linearSequence(16, 0.0, 1.0),
                                                       linearSequence(5, 0.0, 1.0)};
  const std::vector<std::vector<Real>> rewards = {linearSequence(16, 0.0, 1.0),
                                                  linearSequence(5, 0.0, 1.0)};
  const auto trajectories = assembler.assemble({&observations}, {}, {}, rewards);
  ASSERT_EQ(trajectories.size(), 1u);
  EXPECT_EQ(trajectories[0].rewards.size(), 3u);
}
 
TEST(LibtorchRLCoreTest, TrajectoryAssemblerRejectsMismatchedSampleCounts)
{
  const std::vector<std::vector<Real>> first_observations = {linearSequence(3, 0.0, 1.0),
                                                             linearSequence(3, 0.0, 1.0)};
  const std::vector<std::vector<Real>> second_observations = {linearSequence(3, 0.0, 1.0)};
  const std::vector<std::vector<Real>> rewards = {linearSequence(3, 0.0, 1.0),
                                                  linearSequence(3, 0.0, 1.0)};
 
  LibtorchRLTrajectoryAssembler assembler(1, 1, true, false);
  EXPECT_THROW(
      {
        try
        {
          assembler.assemble({&first_observations, &second_observations}, {}, {}, rewards);
        }
        catch (const std::exception & e)
        {
          EXPECT_EQ(std::string(e.what()),
                    "Entry 1 of the observation reporters holds 1 samples, but the reward "
                    "reporter holds 2 samples.");
          throw;
        }
      },
      std::exception);
}
 
TEST(LibtorchRLCoreTest, TrajectoryAssemblerRejectsZeroTimestepWindow)
{
  EXPECT_THROW(
      {
        try
        {
          LibtorchRLTrajectoryAssembler assembler(1, 0, true, false);
          static_cast<void>(assembler);
        }
        catch (const std::exception & e)
        {
          EXPECT_EQ(std::string(e.what()), "The RL trajectory timestep window must be at least 1.");
          throw;
        }
      },
      std::exception);
}
 
TEST(LibtorchRLCoreTest, TrajectoryBufferRewardStatisticsIncludeSpreadBetweenTrajectories)
{
  LibtorchRLTrajectoryBuffer buffer;
  buffer.addTrajectory(trajectoryWithRewards({1.0, 3.0}));
  buffer.addTrajectory(trajectoryWithRewards({5.0}));
 
  const auto statistics = buffer.rewardStatistics();
 
  // The overall mean is 3, so the squared deviations are 4, 0, and 4.
  EXPECT_NEAR(statistics.mean, 3.0, 1e-12);
  EXPECT_NEAR(statistics.std, std::sqrt(8.0 / 3.0), 1e-12);
  EXPECT_EQ(statistics.trajectory_means, std::vector<Real>({2.0, 5.0}));
  EXPECT_EQ(statistics.trajectory_stds, std::vector<Real>({1.0, 0.0}));
 
  buffer.clear();
  const auto empty_statistics = buffer.rewardStatistics();
  EXPECT_EQ(empty_statistics.mean, 0.0);
  EXPECT_EQ(empty_statistics.std, 0.0);
  EXPECT_TRUE(empty_statistics.trajectory_means.empty());
}
 
TEST(LibtorchRLCoreTest, TrajectoryBufferFlattensTrajectoriesInOrder)
{
  const auto options = torch::TensorOptions().dtype(torch::kDouble);
 
  // Values in the tens come from the first trajectory and values in the twenties from the second,
  // so the flattened rows show that trajectories of different lengths are stacked in order.
  auto first = trajectoryWithRewards({10.0, 11.0});
  first.observations = {{10.0, 10.5}, {11.0, 11.5}};
  first.actions = {{-10.0}, {-11.0}};
  first.value_targets = {0.10, 0.11};
  first.advantages = {1.0, 1.1};
  auto second = trajectoryWithRewards({20.0});
  second.observations = {{20.0, 20.5}};
  second.actions = {{-20.0}};
  second.value_targets = {0.20};
  second.advantages = {2.0};
 
  LibtorchRLTrajectoryBuffer buffer;
  buffer.addTrajectory(first);
  buffer.addTrajectory(second);
  EXPECT_EQ(buffer.numTrajectories(), 2u);
  EXPECT_EQ(buffer.numTransitions(), 3u);
 
  const auto batch = buffer.flatten();
  ASSERT_EQ(batch.size(), 3);
  EXPECT_TRUE(torch::equal(batch.observations,
                           torch::tensor({{10.0, 10.5}, {11.0, 11.5}, {20.0, 20.5}}, options)));
  EXPECT_TRUE(torch::equal(batch.next_observations, torch::zeros({3, 1}, options)));
  EXPECT_TRUE(torch::equal(batch.actions, torch::tensor({{-10.0}, {-11.0}, {-20.0}}, options)));
  EXPECT_TRUE(torch::equal(batch.log_probabilities, torch::zeros({3, 1}, options)));
  EXPECT_TRUE(torch::equal(batch.rewards, torch::tensor({{10.0}, {11.0}, {20.0}}, options)));
  EXPECT_TRUE(
      torch::equal(batch.value_targets, torch::tensor({{0.10}, {0.11}, {0.20}}, options)));
  EXPECT_TRUE(torch::equal(batch.advantages, torch::tensor({{1.0}, {1.1}, {2.0}}, options)));
 
  buffer.clear();
  EXPECT_EQ(buffer.flatten().size(), 0);
}
 
TEST(LibtorchRLCoreTest, TrajectoryBufferRejectsMismatchedTrajectoryLengths)
{
  const auto expect_rejected =
      [](const LibtorchRLTrajectoryBuffer::Trajectory & trajectory, const std::string & message)
  {
    LibtorchRLTrajectoryBuffer buffer;
    EXPECT_THROW(
        {
          try
          {
            buffer.addTrajectory(trajectory);
          }
          catch (const std::exception & e)
          {
            EXPECT_EQ(std::string(e.what()), message);
            throw;
          }
        },
        std::exception);
    EXPECT_TRUE(buffer.empty());
  };
 
  // Required fields must match the reward length.
  auto short_actions = trajectoryWithRewards({1.0, 2.0});
  short_actions.actions.pop_back();
  expect_rejected(short_actions, "RL trajectory actions must match the reward sequence length.");
 
  // Optional fields may be empty, but must match the reward length when given.
  auto short_value_targets = trajectoryWithRewards({1.0, 2.0});
  short_value_targets.value_targets = {1.0};
  expect_rejected(short_value_targets,
                  "RL trajectory value targets must match the reward sequence length.");
}
 
} // namespace
 
#endif