Implement automatic observation normalization

Author: Omegastick

example/gym_client.cpp

@@ -16,26 +16,27 @@ using namespace cpprl;
 
 // Algorithm hyperparameters
 const std::string algorithm = "PPO";
-const int batch_size = 2048;
+const float actor_loss_coef = 1.0;
+const int batch_size = 40;
 const float clip_param = 0.2;
 const float discount_factor = 0.99;
-const float entropy_coef = 0.001;
+const float entropy_coef = 1e-3;
 const float gae = 0.95;
-const float learning_rate = 2.5e-4;
+const float kl_target = 0.05;
+const float learning_rate = 7e-4;
 const int log_interval = 1;
 const int max_frames = 10e+7;
-const int num_epoch = 10;
-const int num_mini_batch = 32;
+const int num_epoch = 3;
+const int num_mini_batch = 20;
 const int reward_average_window_size = 10;
+const float reward_clip_value = 10; // Post scaling
 const bool use_gae = true;
-const bool use_lr_decay = true;
-const float actor_loss_coef = 1.0;
+const bool use_lr_decay = false;
 const float value_loss_coef = 0.5;
 
 // Environment hyperparameters
-const float env_gamma = discount_factor; // Set to -1 to disable
-const std::string env_name = "BipedalWalkerHardcore-v2";
-const int num_envs = 16;
+const std::string env_name = "LunarLander-v2";
+const int num_envs = 8;
 const float render_reward_threshold = 160;
 
 // Model hyperparameters
@@ -80,7 +81,6 @@ int main(int argc, char *argv[])
     spdlog::info("Creating environment");
     auto make_param = std::make_shared<MakeParam>();
     make_param->env_name = env_name;
-    make_param->gamma = env_gamma;
     make_param->num_envs = num_envs;
     Request<MakeParam> make_request("make", make_param);
     communicator.send_request(make_request);
@@ -125,7 +125,17 @@ int main(int argc, char *argv[])
     }
     base->to(device);
     ActionSpace space{env_info->action_space_type, env_info->action_space_shape};
-    Policy policy(space, base);
+    Policy policy(nullptr);
+    if (env_info->observation_space_shape.size() == 1)
+    {
+        // With observation normalization
+        policy = Policy(space, base, true);
+    }
+    else
+    {
+        // Without observation normalization
+        policy = Policy(space, base, true);
+    }
     policy->to(device);
     RolloutStorage storage(batch_size, num_envs, env_info->observation_space_shape, space, hidden_size, device);
     std::unique_ptr<Algorithm> algo;
@@ -135,7 +145,17 @@ int main(int argc, char *argv[])
     }
     else if (algorithm == "PPO")
     {
-        algo = std::make_unique<PPO>(policy, clip_param, num_epoch, num_mini_batch, actor_loss_coef, value_loss_coef, entropy_coef, learning_rate);
+        algo = std::make_unique<PPO>(policy,
+                                     clip_param,
+                                     num_epoch,
+                                     num_mini_batch,
+                                     actor_loss_coef,
+                                     value_loss_coef,
+                                     entropy_coef,
+                                     learning_rate,
+                                     1e-8,
+                                     0.5,
+                                     kl_target);
     }
 
     storage.set_first_observation(observation);
@@ -144,6 +164,8 @@ int main(int argc, char *argv[])
     int episode_count = 0;
     bool render = false;
     std::vector<float> reward_history(reward_average_window_size);
+    RunningMeanStd returns_rms(1);
+    auto returns = torch::zeros({num_envs});
 
     auto start_time = std::chrono::high_resolution_clock::now();
 
@@ -159,14 +181,21 @@ int main(int argc, char *argv[])
                                          storage.get_hidden_states()[step],
                                          storage.get_masks()[step]);
             }
-            auto actions_tensor = act_result[1].cpu();
+            auto actions_tensor = act_result[1].cpu().to(torch::kFloat);
             float *actions_array = actions_tensor.data<float>();
             std::vector<std::vector<float>> actions(num_envs);
             for (int i = 0; i < num_envs; ++i)
             {
-                for (int j = 0; j < env_info->action_space_shape[0]; j++)
+                if (space.type == "Discrete")
+                {
+                    actions[i] = {actions_array[i]};
+                }
+                else
                 {
-                    actions[i].push_back(actions_array[i * env_info->action_space_shape[0] + j]);
+                    for (int j = 0; j < env_info->action_space_shape[0]; j++)
+                    {
+                        actions[i].push_back(actions_array[i * env_info->action_space_shape[0] + j]);
+                    }
                 }
             }
 
@@ -183,7 +212,13 @@ int main(int argc, char *argv[])
                 auto step_result = communicator.get_response<CnnStepResponse>();
                 observation_vec = flatten_vector(step_result->observation);
                 observation = torch::from_blob(observation_vec.data(), observation_shape).to(device);
-                rewards = flatten_vector(step_result->reward);
+                auto raw_reward_vec = flatten_vector(step_result->real_reward);
+                auto reward_tensor = torch::from_blob(raw_reward_vec.data(), {num_envs}, torch::kFloat);
+                returns = returns * discount_factor + reward_tensor;
+                returns_rms->update(returns);
+                reward_tensor = torch::clamp(reward_tensor / torch::sqrt(returns_rms->get_variance() + 1e-8),
+                                             -reward_clip_value, reward_clip_value);
+                rewards = std::vector<float>(reward_tensor.data<float>(), reward_tensor.data<float>() + reward_tensor.numel());
                 real_rewards = flatten_vector(step_result->real_reward);
                 dones_vec = step_result->done;
             }
@@ -192,7 +227,13 @@ int main(int argc, char *argv[])
                 auto step_result = communicator.get_response<MlpStepResponse>();
                 observation_vec = flatten_vector(step_result->observation);
                 observation = torch::from_blob(observation_vec.data(), observation_shape).to(device);
-                rewards = flatten_vector(step_result->reward);
+                auto raw_reward_vec = flatten_vector(step_result->real_reward);
+                auto reward_tensor = torch::from_blob(raw_reward_vec.data(), {num_envs}, torch::kFloat);
+                returns = returns * discount_factor + reward_tensor;
+                returns_rms->update(returns);
+                reward_tensor = torch::clamp(reward_tensor / torch::sqrt(returns_rms->get_variance() + 1e-8),
+                                             -reward_clip_value, reward_clip_value);
+                rewards = std::vector<float>(reward_tensor.data<float>(), reward_tensor.data<float>() + reward_tensor.numel());
                 real_rewards = flatten_vector(step_result->real_reward);
                 dones_vec = step_result->done;
             }
@@ -203,6 +244,7 @@ int main(int argc, char *argv[])
                 {
                     reward_history[episode_count % reward_average_window_size] = running_rewards[i];
                     running_rewards[i] = 0;
+                    returns[i] = 0;
                     episode_count++;
                 }
             }

example/requests.h

@@ -25,9 +25,8 @@ struct InfoParam
 struct MakeParam
 {
     std::string env_name;
-    float gamma;
     int num_envs;
-    MSGPACK_DEFINE_MAP(env_name, gamma, num_envs);
+    MSGPACK_DEFINE_MAP(env_name, num_envs);
 };
 
 struct ResetParam

gym_server/envs.py

@@ -138,21 +138,15 @@ def _thunk():
     return _thunk
 
 
-def make_vec_envs(env_name, seed, num_processes, gamma, num_frame_stack=None):
+def make_vec_envs(env_name, seed, num_processes, num_frame_stack=None):
     envs = [make_env(env_name, seed, i) for i in range(num_processes)]
 
     if len(envs) > 1:
         envs = SubprocVecEnv(envs)
     else:
         envs = DummyVecEnv(envs)
 
-    if len(envs.observation_space.shape) == 1:
-        if gamma is None or gamma == -1:
-            envs = VecNormalize(envs, ret=False)
-        else:
-            envs = VecNormalize(envs, gamma=gamma)
-    else:
-        envs = VecRewardInfo(envs)
+    envs = VecRewardInfo(envs)
 
     if num_frame_stack is not None:
         envs = VecFrameStack(envs, num_frame_stack)

gym_server/server.py

@@ -53,8 +53,7 @@ def _serve(self):
                                                  observation_space_shape))
 
             elif method == 'make':
-                self.__make(param['env_name'], param['num_envs'],
-                            param['gamma'])
+                self.__make(param['env_name'], param['num_envs'])
                 self.zmq_client.send(MakeMessage())
 
             elif method == 'reset':
@@ -86,12 +85,12 @@ def info(self):
         return (action_space_type, action_space_shape, observation_space_type,
                 observation_space_shape)
 
-    def make(self, env_name, num_envs, gamma):
+    def make(self, env_name, num_envs):
         """
         Makes a vectorized environment of the type and number specified.
         """
         logging.info("Making %d %ss", num_envs, env_name)
-        self.env = make_vec_envs(env_name, 0, num_envs, gamma)
+        self.env = make_vec_envs(env_name, 0, num_envs)
 
     def reset(self) -> np.ndarray:
         """
@@ -109,6 +108,7 @@ def step(self,
         """
         if isinstance(self.env.action_space, gym.spaces.Discrete):
             actions = actions.squeeze(-1)
+            actions = actions.astype(np.int)
         observation, reward, done, info = self.env.step(actions)
         reward = np.expand_dims(reward, -1)
         done = np.expand_dims(done, -1)

include/cpprl/cpprl.h

@@ -11,5 +11,6 @@
 #include "cpprl/model/nn_base.h"
 #include "cpprl/model/output_layers.h"
 #include "cpprl/model/policy.h"
+#include "cpprl/observation_normalizer.h"
 #include "cpprl/spaces.h"
 #include "cpprl/storage.h"

include/cpprl/model/mlp_base.h

@@ -16,6 +16,7 @@ class MlpBase : public NNBase
     nn::Sequential actor;
     nn::Sequential critic;
     nn::Linear critic_linear;
+    unsigned int num_inputs;
 
   public:
     MlpBase(unsigned int num_inputs,
@@ -25,5 +26,7 @@ class MlpBase : public NNBase
     std::vector<torch::Tensor> forward(torch::Tensor inputs,
                                        torch::Tensor hxs,
                                        torch::Tensor masks);
+
+    inline unsigned int get_num_inputs() const { return num_inputs; }
 };
 }

include/cpprl/model/nn_base.h

@@ -11,9 +11,9 @@ namespace cpprl
 class NNBase : public nn::Module
 {
   private:
-    bool recurrent;
-    unsigned int hidden_size;
     nn::GRU gru;
+    unsigned int hidden_size;
+    bool recurrent;
 
   public:
     NNBase(bool recurrent,

include/cpprl/model/policy.h

@@ -7,6 +7,7 @@
 
 #include "cpprl/model/nn_base.h"
 #include "cpprl/model/output_layers.h"
+#include "cpprl/observation_normalizer.h"
 #include "cpprl/spaces.h"
 
 using namespace torch;
@@ -16,16 +17,19 @@ namespace cpprl
 class PolicyImpl : public nn::Module
 {
   private:
+    ActionSpace action_space;
     std::shared_ptr<NNBase> base;
+    ObservationNormalizer observation_normalizer;
     std::shared_ptr<OutputLayer> output_layer;
-    ActionSpace action_space;
 
     std::vector<torch::Tensor> forward_gru(torch::Tensor x,
                                            torch::Tensor hxs,
                                            torch::Tensor masks);
 
   public:
-    PolicyImpl(ActionSpace action_space, std::shared_ptr<NNBase> base);
+    PolicyImpl(ActionSpace action_space,
+               std::shared_ptr<NNBase> base,
+               bool normalize_observations = false);
 
     std::vector<torch::Tensor> act(torch::Tensor inputs,
                                    torch::Tensor rnn_hxs,
@@ -40,12 +44,17 @@ class PolicyImpl : public nn::Module
     torch::Tensor get_values(torch::Tensor inputs,
                              torch::Tensor rnn_hxs,
                              torch::Tensor masks);
+    void update_observation_normalizer(torch::Tensor observations);
 
     inline bool is_recurrent() const { return base->is_recurrent(); }
     inline unsigned int get_hidden_size() const
     {
         return base->get_hidden_size();
     }
+    inline bool using_observation_normalizer() const
+    {
+        return !observation_normalizer.is_empty();
+    }
 };
 TORCH_MODULE(Policy);
 }

include/cpprl/observation_normalizer.h

@@ -6,7 +6,7 @@
 
 #include "cpprl/running_mean_std.h"
 
-namespace SingularityTrainer
+namespace cpprl
 {
 class ObservationNormalizer;
 

include/cpprl/running_mean_std.h

@@ -4,7 +4,7 @@
 
 #include <torch/torch.h>
 
-namespace SingularityTrainer
+namespace cpprl
 {
 // https://github.com/openai/baselines/blob/master/baselines/common/running_mean_std.py
 class RunningMeanStdImpl : public torch::nn::Module