slight modularization

src/csrc/include/internal/rl/off_policy.h

@@ -57,7 +57,7 @@ class RLOffPolicySystem {
  RLOffPolicySystem(const RLOffPolicySystem&) = delete;
 
  // empty constructor:
- RLOffPolicySystem() : train_step_count_(0) {}
+ RLOffPolicySystem(int model_device, int rb_device);
 
  // some important functions which have to be implemented by the base class
  virtual void updateReplayBuffer(torch::Tensor, torch::Tensor, torch::Tensor, float, bool) = 0;
@@ -80,6 +80,7 @@ class RLOffPolicySystem {
  virtual std::shared_ptr<ModelState> getSystemState_() = 0;
  virtual std::shared_ptr<Comm> getSystemComm_() = 0;
  size_t train_step_count_;
+ torch::Device model_device_, rb_device_;
 };
 
 // Declaration of external global variables

src/csrc/include/internal/rl/on_policy.h

@@ -56,8 +56,8 @@ class RLOnPolicySystem {
  // disable copy constructor
  RLOnPolicySystem(const RLOnPolicySystem&) = delete;
 
- // empty constructor:
- RLOnPolicySystem() : train_step_count_(0) {}
+ // default constructor:
+ RLOnPolicySystem(int model_device, int rb_device);
 
  // some important functions which have to be implemented by the base class
  virtual void updateRolloutBuffer(torch::Tensor, torch::Tensor, float, float, float, bool) = 0;
@@ -81,6 +81,7 @@ class RLOnPolicySystem {
  virtual std::shared_ptr<ModelState> getSystemState_() = 0;
  virtual std::shared_ptr<Comm> getSystemComm_() = 0;
  size_t train_step_count_;
+ torch::Device model_device_, rb_device_;
 };
 
 // Declaration of external global variables

src/csrc/include/internal/rl/on_policy/ppo.h

@@ -277,10 +277,6 @@ class PPOSystem : public RLOnPolicySystem, public std::enable_shared_from_this<R
 
  std::shared_ptr<Comm> getSystemComm_();
 
- // device
- torch::Device model_device_;
- torch::Device rb_device_;
-
  // models
  ACPolicyPack p_model_;
  ModelPack q_model_;

src/csrc/rl/off_policy/interface.cpp

@@ -68,6 +68,7 @@ RLOffPolicySystem::RLOffPolicySystem(int model_device, int rb_device) : train_st
 
 } // namespace off_policy
 
+} // namespace rl
 } // namespace torchfort
 
 torchfort_result_t torchfort_rl_off_policy_create_system(const char* name, const char* config_fname,

src/csrc/rl/on_policy/interface.cpp

@@ -56,6 +56,14 @@ namespace rl {
 namespace on_policy {
 // Global variables
 std::unordered_map<std::string, std::shared_ptr<RLOnPolicySystem>> registry;
+
+// default constructor:
+RLOnPolicySystem::RLOnPolicySystem(int model_device, int rb_device) : train_step_count_(0), model_device_(get_device(model_device)), rb_device_(get_device(rb_device)) {
+ if ( !(torchfort::rl::validate_devices(model_device, rb_device)) ) {
+ THROW_INVALID_USAGE("The parameters model_device and rb_device have to specify the same GPU or one has to specify a GPU and the other the CPU.");
+ }
+}
+
 } // namespace on_policy
 
 } // namespace rl

src/csrc/rl/on_policy/ppo.cpp

@@ -43,7 +43,7 @@ namespace on_policy {
 
 PPOSystem::PPOSystem(const char* name, const YAML::Node& system_node,
  int model_device, int rb_device)
- : model_device_(get_device(model_device)), rb_device_(get_device(rb_device)) {
+ : RLOnPolicySystem(model_device, rb_device) {
 
  // get basic parameters first
  auto algo_node = system_node["algorithm"];