Fix global log_std implementation

Author: CovERUshKA

src/engine/server/NN/ModelManager.cpp

@@ -26,7 +26,7 @@ int64_t h_start = 1024; // 1024 256
 double std_dev = 0.37; // log(0.37) ~ -1
 double learning_rate = 5e-5;
 double actor_learning_rate = 3e-4;
-//double log_std_learning_rate = 1e-4;
+double log_std_learning_rate = 1e-4;
 double critic_learning_rate = 1e-3;
 //double weight_decay = 0.0001;
 
@@ -158,8 +158,8 @@ ModelManager::ModelManager(bool is_training, std::string train_folder, size_t ba
 							std::make_unique<torch::optim::AdamOptions>(actor_learning_rate)));
 	param_groups.push_back(torch::optim::OptimizerParamGroup({ac_update->critic_network->parameters()},
 							std::make_unique<torch::optim::AdamOptions>(critic_learning_rate)));
-	//param_groups.push_back(torch::optim::OptimizerParamGroup({ac_update->log_std_},
-		//std::make_unique<torch::optim::AdamOptions>(log_std_learning_rate)));
+	param_groups.push_back(torch::optim::OptimizerParamGroup({ac_update->log_std_},
+			std::make_unique<torch::optim::AdamOptions>(log_std_learning_rate)));
 
 	opt = std::make_shared<torch::optim::Adam>(param_groups);
 

src/engine/server/NN/Models.h

@@ -234,7 +234,7 @@ struct ActorCriticImpl : public torch::nn::Module
 	    auto gaussian_entropy = 0.5 + 0.5 * log(2 * M_PI) + log_std_;
 	    //std::cout << gaussian_entropy.sizes() << std::endl;
 	    // Sum over the last dimension (angle components)
-	    return gaussian_entropy.sum(1); // Shape [...]
+	    return gaussian_entropy.sum(); // Shape [...]
     }
 
 	// Bernoulli entropy