Convert ObservationNormalizer and RunningMeanStd to Torch modules

Omegastick · Omegastick · commit 0615f7646715 · 2019-07-23T13:10:00.000+09:00
diff --git a/include/cpprl/observation_normalizer.h b/include/cpprl/observation_normalizer.h
@@ -8,25 +8,28 @@
 
 namespace SingularityTrainer
 {
-class ObservationNormalizer
+class ObservationNormalizer;
+
+class ObservationNormalizerImpl : public torch::nn::Module
 {
   private:
-    float clip;
+    torch::Tensor clip;
     RunningMeanStd rms;
 
   public:
-    explicit ObservationNormalizer(int size, float clip = 10.);
-    ObservationNormalizer(const std::vector<float> &means,
-                          const std::vector<float> &variances,
-                          float clip = 10.);
-    explicit ObservationNormalizer(const std::vector<ObservationNormalizer> &others);
+    explicit ObservationNormalizerImpl(int size, float clip = 10.);
+    ObservationNormalizerImpl(const std::vector<float> &means,
+                              const std::vector<float> &variances,
+                              float clip = 10.);
+    explicit ObservationNormalizerImpl(const std::vector<ObservationNormalizer> &others);
 
     torch::Tensor process_observation(torch::Tensor observation);
     std::vector<float> get_mean() const;
     std::vector<float> get_variance() const;
     void update(torch::Tensor observations);
 
-    inline float get_clip_value() const { return clip; }
-    inline int get_step_count() const { return rms.get_count(); }
+    inline float get_clip_value() const { return clip.item().toFloat(); }
+    inline int get_step_count() const { return rms->get_count(); }
 };
+TORCH_MODULE(ObservationNormalizer);
 }
diff --git a/include/cpprl/running_mean_std.h b/include/cpprl/running_mean_std.h
@@ -7,25 +7,25 @@
 namespace SingularityTrainer
 {
 // https://github.com/openai/baselines/blob/master/baselines/common/running_mean_std.py
-class RunningMeanStd
+class RunningMeanStdImpl : public torch::nn::Module
 {
   private:
-    float count;
-    torch::Tensor mean, variance;
+    torch::Tensor count, mean, variance;
 
     void update_from_moments(torch::Tensor batch_mean,
                              torch::Tensor batch_var,
                              int batch_count);
 
   public:
-    explicit RunningMeanStd(int size);
-    RunningMeanStd(std::vector<float> means, std::vector<float> variances);
+    explicit RunningMeanStdImpl(int size);
+    RunningMeanStdImpl(std::vector<float> means, std::vector<float> variances);
 
     void update(torch::Tensor observation);
 
-    inline int get_count() const { return static_cast<int>(count); }
+    inline int get_count() const { return static_cast<int>(count.item().toFloat()); }
     inline torch::Tensor get_mean() const { return mean.clone(); }
     inline torch::Tensor get_variance() const { return variance.clone(); }
-    inline void set_count(int count) { this->count = count + 1e-8; }
+    inline void set_count(int count) { this->count[0] = count + 1e-8; }
 };
+TORCH_MODULE(RunningMeanStd);
 }
diff --git a/src/observation_normalizer.cpp b/src/observation_normalizer.cpp
@@ -6,32 +6,32 @@
 
 namespace SingularityTrainer
 {
-ObservationNormalizer::ObservationNormalizer(int size, float clip)
-    : clip(clip),
-      rms(size) {}
-
-ObservationNormalizer::ObservationNormalizer(const std::vector<float> &means,
-                                             const std::vector<float> &variances,
-                                             float clip)
-    : clip(clip),
-      rms(means, variances){}
-
-ObservationNormalizer::ObservationNormalizer(const std::vector<ObservationNormalizer> &others)
-    : clip(0),
-      rms(1)
+ObservationNormalizerImpl::ObservationNormalizerImpl(int size, float clip)
+    : clip(register_buffer("clip", torch::full({1}, clip, torch::kFloat))),
+      rms(register_module("rms", RunningMeanStd(size))) {}
+
+ObservationNormalizerImpl::ObservationNormalizerImpl(const std::vector<float> &means,
+                                                     const std::vector<float> &variances,
+                                                     float clip)
+    : clip(register_buffer("clip", torch::full({1}, clip, torch::kFloat))),
+      rms(register_module("rms", RunningMeanStd(means, variances))) {}
+
+ObservationNormalizerImpl::ObservationNormalizerImpl(const std::vector<ObservationNormalizer> &others)
+    : clip(register_buffer("clip", torch::zeros({1}, torch::kFloat))),
+      rms(register_module("rms", RunningMeanStd(1)))
 {
     // Calculate mean clip
     for (const auto &other : others)
     {
-        clip += other.get_clip_value();
+        clip += other->get_clip_value();
     }
-    clip /= others.size();
+    clip[0] = clip[0] / static_cast<float>(others.size());
 
     // Calculate mean mean
-    std::vector<float> mean_means(others[0].get_mean().size(), 0);
+    std::vector<float> mean_means(others[0]->get_mean().size(), 0);
     for (const auto &other : others)
     {
-        auto other_mean = other.get_mean();
+        auto other_mean = other->get_mean();
         for (unsigned int i = 0; i < mean_means.size(); ++i)
         {
             mean_means[i] += other_mean[i];
@@ -43,10 +43,10 @@ ObservationNormalizer::ObservationNormalizer(const std::vector<ObservationNormal
     }
 
     // Calculate mean variances
-    std::vector<float> mean_variances(others[0].get_variance().size(), 0);
+    std::vector<float> mean_variances(others[0]->get_variance().size(), 0);
     for (const auto &other : others)
     {
-        auto other_variances = other.get_variance();
+        auto other_variances = other->get_variance();
         for (unsigned int i = 0; i < mean_variances.size(); ++i)
         {
             mean_variances[i] += other_variances[i];
@@ -61,33 +61,33 @@ ObservationNormalizer::ObservationNormalizer(const std::vector<ObservationNormal
 
     int total_count = std::accumulate(others.begin(), others.end(), 0,
                                       [](int accumulator, const ObservationNormalizer &other) {
-                                          return accumulator + other.get_step_count();
+                                          return accumulator + other->get_step_count();
                                       });
-    rms.set_count(total_count);
+    rms->set_count(total_count);
 }
 
-torch::Tensor ObservationNormalizer::process_observation(torch::Tensor observation)
+torch::Tensor ObservationNormalizerImpl::process_observation(torch::Tensor observation)
 {
-    auto normalized_obs = (observation - rms.get_mean()) /
-                          torch::sqrt(rms.get_variance() + 1e-8);
-    return torch::clamp(normalized_obs, -clip, clip);
+    auto normalized_obs = (observation - rms->get_mean()) /
+                          torch::sqrt(rms->get_variance() + 1e-8);
+    return torch::clamp(normalized_obs, -clip.item(), clip.item());
 }
 
-std::vector<float> ObservationNormalizer::get_mean() const
+std::vector<float> ObservationNormalizerImpl::get_mean() const
 {
-    auto mean = rms.get_mean();
+    auto mean = rms->get_mean();
     return std::vector<float>(mean.data<float>(), mean.data<float>() + mean.numel());
 }
 
-std::vector<float> ObservationNormalizer::get_variance() const
+std::vector<float> ObservationNormalizerImpl::get_variance() const
 {
-    auto variance = rms.get_variance();
+    auto variance = rms->get_variance();
     return std::vector<float>(variance.data<float>(), variance.data<float>() + variance.numel());
 }
 
-void ObservationNormalizer::update(torch::Tensor observations)
+void ObservationNormalizerImpl::update(torch::Tensor observations)
 {
-    rms.update(observations);
+    rms->update(observations);
 }
 
 TEST_CASE("ObservationNormalizer")
@@ -97,7 +97,7 @@ TEST_CASE("ObservationNormalizer")
         ObservationNormalizer normalizer(7, 1);
         float observation_array[] = {-1000, -100, -10, 0, 10, 100, 1000};
         auto observation = torch::from_blob(observation_array, {7});
-        auto processed_observation = normalizer.process_observation(observation);
+        auto processed_observation = normalizer->process_observation(observation);
 
         auto has_too_large_values = (processed_observation > 1).any().item().toBool();
         auto has_too_small_values = (processed_observation < -1).any().item().toBool();
@@ -116,10 +116,10 @@ TEST_CASE("ObservationNormalizer")
         auto obs_2 = torch::from_blob(obs_2_array, {5});
         auto obs_3 = torch::from_blob(obs_3_array, {5});
 
-        normalizer.update(obs_1);
-        normalizer.update(obs_2);
-        normalizer.update(obs_3);
-        auto processed_observation = normalizer.process_observation(obs_3);
+        normalizer->update(obs_1);
+        normalizer->update(obs_2);
+        normalizer->update(obs_3);
+        auto processed_observation = normalizer->process_observation(obs_3);
 
         DOCTEST_CHECK(processed_observation[0].item().toFloat() == doctest::Approx(1.26008659));
         DOCTEST_CHECK(processed_observation[1].item().toFloat() == doctest::Approx(0.70712887));
@@ -130,10 +130,10 @@ TEST_CASE("ObservationNormalizer")
 
     SUBCASE("Loads mean and variance from constructor correctly")
     {
-        ObservationNormalizer normalizer({1, 2, 3}, {4, 5, 6});
+        ObservationNormalizer normalizer(std::vector<float>({1, 2, 3}), std::vector<float>({4, 5, 6}));
 
-        auto mean = normalizer.get_mean();
-        auto variance = normalizer.get_variance();
+        auto mean = normalizer->get_mean();
+        auto variance = normalizer->get_variance();
         DOCTEST_CHECK(mean[0] == doctest::Approx(1));
         DOCTEST_CHECK(mean[1] == doctest::Approx(2));
         DOCTEST_CHECK(mean[2] == doctest::Approx(3));
@@ -150,18 +150,18 @@ TEST_CASE("ObservationNormalizer")
             normalizers.push_back(ObservationNormalizer(3));
             for (int j = 0; j <= i; ++j)
             {
-                normalizers[i].update(torch::rand({3}));
+                normalizers[i]->update(torch::rand({3}));
             }
         }
 
         ObservationNormalizer combined_normalizer(normalizers);
 
         std::vector<std::vector<float>> means;
         std::transform(normalizers.begin(), normalizers.end(), std::back_inserter(means),
-                       [](const ObservationNormalizer &normalizer) { return normalizer.get_mean(); });
+                       [](const ObservationNormalizer &normalizer) { return normalizer->get_mean(); });
         std::vector<std::vector<float>> variances;
         std::transform(normalizers.begin(), normalizers.end(), std::back_inserter(variances),
-                       [](const ObservationNormalizer &normalizer) { return normalizer.get_variance(); });
+                       [](const ObservationNormalizer &normalizer) { return normalizer->get_variance(); });
 
         std::vector<float> mean_means;
         for (int i = 0; i < 3; ++i)
@@ -174,15 +174,15 @@ TEST_CASE("ObservationNormalizer")
             mean_variances.push_back((variances[0][i] + variances[1][i] + variances[2][i]) / 3);
         }
 
-        auto actual_mean_means = combined_normalizer.get_mean();
-        auto actual_mean_variances = combined_normalizer.get_variance();
+        auto actual_mean_means = combined_normalizer->get_mean();
+        auto actual_mean_variances = combined_normalizer->get_variance();
 
         for (int i = 0; i < 3; ++i)
         {
             DOCTEST_CHECK(actual_mean_means[i] == doctest::Approx(mean_means[i]));
             DOCTEST_CHECK(actual_mean_variances[i] == doctest::Approx(actual_mean_variances[i]));
         }
-        DOCTEST_CHECK(combined_normalizer.get_step_count() == 6);
+        DOCTEST_CHECK(combined_normalizer->get_step_count() == 6);
     }
 }
 }
diff --git a/src/running_mean_std.cpp b/src/running_mean_std.cpp
@@ -5,19 +5,19 @@
 
 namespace SingularityTrainer
 {
-RunningMeanStd::RunningMeanStd(int size)
-    : count(1e-4),
-      mean(torch::zeros({size})),
-      variance(torch::ones({size})) {}
+RunningMeanStdImpl::RunningMeanStdImpl(int size)
+    : count(register_buffer("count", torch::full({1}, 1e-4, torch::kFloat))),
+      mean(register_buffer("mean", torch::zeros({size}))),
+      variance(register_buffer("variance", torch::ones({size}))) {}
 
-RunningMeanStd::RunningMeanStd(std::vector<float> means, std::vector<float> variances)
-    : count(1e-4),
-      mean(torch::from_blob(means.data(), {static_cast<long>(means.size())})
-               .clone()),
-      variance(torch::from_blob(variances.data(), {static_cast<long>(variances.size())})
-                   .clone()) {}
+RunningMeanStdImpl::RunningMeanStdImpl(std::vector<float> means, std::vector<float> variances)
+    : count(register_buffer("count", torch::full({1}, 1e-4, torch::kFloat))),
+      mean(register_buffer("mean", torch::from_blob(means.data(), {static_cast<long>(means.size())})
+                                       .clone())),
+      variance(register_buffer("variance", torch::from_blob(variances.data(), {static_cast<long>(variances.size())})
+                                               .clone())) {}
 
-void RunningMeanStd::update(torch::Tensor observation)
+void RunningMeanStdImpl::update(torch::Tensor observation)
 {
     observation = observation.reshape({-1, mean.size(0)});
     auto batch_mean = observation.mean(0);
@@ -27,12 +27,12 @@ void RunningMeanStd::update(torch::Tensor observation)
     update_from_moments(batch_mean, batch_var, batch_count);
 }
 
-void RunningMeanStd::update_from_moments(torch::Tensor batch_mean,
-                                         torch::Tensor batch_var,
-                                         int batch_count)
+void RunningMeanStdImpl::update_from_moments(torch::Tensor batch_mean,
+                                             torch::Tensor batch_var,
+                                             int batch_count)
 {
     auto delta = batch_mean - mean;
-    float total_count = count + batch_count;
+    auto total_count = count + batch_count;
 
     mean = mean + delta * batch_count / total_count;
     auto m_a = variance * count;
@@ -48,15 +48,15 @@ TEST_CASE("RunningMeanStd")
     {
         RunningMeanStd rms(5);
         auto observations = torch::rand({3, 5});
-        rms.update(observations[0]);
-        rms.update(observations[1]);
-        rms.update(observations[2]);
+        rms->update(observations[0]);
+        rms->update(observations[1]);
+        rms->update(observations[2]);
 
         auto expected_mean = observations.mean(0);
         auto expected_variance = observations.var(0, false, false);
 
-        auto actual_mean = rms.get_mean();
-        auto actual_variance = rms.get_variance();
+        auto actual_mean = rms->get_mean();
+        auto actual_variance = rms->get_variance();
 
         for (int i = 0; i < 5; ++i)
         {
@@ -71,10 +71,10 @@ TEST_CASE("RunningMeanStd")
 
     SUBCASE("Loads mean and variance from constructor correctly")
     {
-        RunningMeanStd rms({1, 2, 3}, {4, 5, 6});
+        RunningMeanStd rms(std::vector<float>{1, 2, 3}, std::vector<float>{4, 5, 6});
 
-        auto mean = rms.get_mean();
-        auto variance = rms.get_variance();
+        auto mean = rms->get_mean();
+        auto variance = rms->get_variance();
         DOCTEST_CHECK(mean[0].item().toFloat() == doctest::Approx(1));
         DOCTEST_CHECK(mean[1].item().toFloat() == doctest::Approx(2));
         DOCTEST_CHECK(mean[2].item().toFloat() == doctest::Approx(3));
