[Docs] Use more appropriate ActorValueOperator in PPOLoss documentati…

…on (pytorch#2350)
ai4co · Aug 2, 2024 · 59d2ae1 · 59d2ae1
1 parent 99332f5
commit 59d2ae1
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/torchrl/objectives/ppo.py b/torchrl/objectives/ppo.py
@@ -141,7 +141,7 @@ class PPOLoss(LossModule):
         >>> actor_head = SomeActor(in_keys=["hidden"])
         >>> value_head = SomeValue(in_keys=["hidden"])
         >>> # first option, with 2 calls on the common module
-        >>> model = ActorCriticOperator(common, actor_head, value_head)
+        >>> model = ActorValueOperator(common, actor_head, value_head)
         >>> loss_module = PPOLoss(model.get_policy_operator(), model.get_value_operator())
         >>> # second option, with a single call to the common module
         >>> loss_module = PPOLoss(ProbabilisticTensorDictSequential(model, actor_head), value_head)
@@ -718,10 +718,10 @@ class ClipPPOLoss(PPOLoss):
         >>> actor_head = SomeActor(in_keys=["hidden"])
         >>> value_head = SomeValue(in_keys=["hidden"])
         >>> # first option, with 2 calls on the common module
-        >>> model = ActorCriticOperator(common, actor_head, value_head)
-        >>> loss_module = PPOLoss(model.get_policy_operator(), model.get_value_operator())
+        >>> model = ActorValueOperator(common, actor_head, value_head)
+        >>> loss_module = ClipPPOLoss(model.get_policy_operator(), model.get_value_operator())
         >>> # second option, with a single call to the common module
-        >>> loss_module = PPOLoss(ProbabilisticTensorDictSequential(model, actor_head), value_head)
+        >>> loss_module = ClipPPOLoss(ProbabilisticTensorDictSequential(model, actor_head), value_head)
 
       This will work regardless of whether separate_losses is activated or not.
 
@@ -955,10 +955,10 @@ class KLPENPPOLoss(PPOLoss):
         >>> actor_head = SomeActor(in_keys=["hidden"])
         >>> value_head = SomeValue(in_keys=["hidden"])
         >>> # first option, with 2 calls on the common module
-        >>> model = ActorCriticOperator(common, actor_head, value_head)
-        >>> loss_module = PPOLoss(model.get_policy_operator(), model.get_value_operator())
+        >>> model = ActorValueOperator(common, actor_head, value_head)
+        >>> loss_module = KLPENPPOLoss(model.get_policy_operator(), model.get_value_operator())
         >>> # second option, with a single call to the common module
-        >>> loss_module = PPOLoss(ProbabilisticTensorDictSequential(model, actor_head), value_head)
+        >>> loss_module = KLPENPPOLoss(ProbabilisticTensorDictSequential(model, actor_head), value_head)
 
       This will work regardless of whether separate_losses is activated or not.