wip

Signed-off-by: sven1977 <[email protected]>
ray-project · Oct 3, 2024 · 09bac52 · 09bac52
1 parent 3bf0d09
commit 09bac52
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/rllib/examples/learners/classes/vpg_learner.py b/rllib/examples/learners/classes/vpg_learner.py
@@ -23,9 +23,7 @@ def build(self) -> None:
         # Prepend the returns-to-go connector piece to have that information
         # available in the train batch.
         if self.config.add_default_connectors_to_learner_pipeline:
-            self._learner_connector.prepend(
-                ComputeReturnsToGo(gamma=self.config.gamma)
-            )
+            self._learner_connector.prepend(ComputeReturnsToGo(gamma=self.config.gamma))
 
     @override(TorchLearner)
     def compute_loss_for_module(

diff --git a/rllib/examples/rl_modules/classes/vpg_rlm.py b/rllib/examples/rl_modules/classes/vpg_rlm.py
@@ -18,7 +18,7 @@ def setup(self):
         hidden_dim = self.model_config["hidden_dim"]
         output_dim = self.action_space.n
 
-        self.policy = nn.Sequential(
+        self._policy_net = nn.Sequential(
             nn.Linear(input_dim, hidden_dim),
             nn.ReLU(),
             nn.Linear(hidden_dim, output_dim),
@@ -36,5 +36,5 @@ def _forward_exploration(self, batch, **kwargs):
 
     @override(TorchRLModule)
     def _forward_train(self, batch, **kwargs):
-        action_logits = self.policy(batch[Columns.OBS])
+        action_logits = self._policy_net(batch[Columns.OBS])
         return {Columns.ACTION_DIST_INPUTS: action_logits}