amend

pytorch · Mar 19, 2024 · 7de1582 · 7de1582
1 parent db63217
commit 7de1582
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/examples/agents/multi-step.py b/examples/agents/multi-step.py
@@ -10,6 +10,9 @@
 
 The core component of this example is the `MultiStepActorWrapper` class.
 
+`MultiStepActorWrapper` handles the calls to the actor when the macro has run out of actions or
+when the environment has been reset (which is indicated by the InitTracker transform).
+
 """
 
 import torch.nn
@@ -22,13 +25,14 @@
     SerialEnv,
     TransformedEnv,
 )
-from torchrl.modules.tensordict_module.actors import Actor, MultiStepActorWrapper
+from torchrl.modules.tensordict_module.actors import MultiStepActorWrapper
 
 time_steps = 6
 n_obs = 4
 n_action = 2
 batch = 5
 
+
 # Transforms a CatFrames in a stack of frames
 def reshape_cat(data: torch.Tensor):
     return data.unflatten(-1, (time_steps, n_obs))