functional fixes

Author: ClemensSchwarke

rsl_rl/algorithms/distillation.py

@@ -53,7 +53,7 @@ def __init__(
 
         # Initialize the transition
         self.transition = RolloutStorage.Transition()
-        self.last_hidden_states = None
+        self.last_hidden_states = (None, None)
 
         # Distillation parameters
         self.num_learning_epochs = num_learning_epochs

rsl_rl/networks/memory.py

@@ -20,7 +20,7 @@ class Memory(nn.Module):
     def __init__(self, input_size: int, hidden_dim: int = 256, num_layers: int = 1, type: str = "lstm") -> None:
         super().__init__()
         rnn_cls = nn.GRU if type.lower() == "gru" else nn.LSTM
-        self.rnn = rnn_cls(input_size=input_size, hidden_dim=hidden_dim, num_layers=num_layers)
+        self.rnn = rnn_cls(input_size=input_size, hidden_size=hidden_dim, num_layers=num_layers)
         self.hidden_states = None
 
     def forward(

rsl_rl/runners/__init__.py

@@ -5,7 +5,7 @@
 
 """Implementation of runners for environment-agent interaction."""
 
+from .on_policy_runner import OnPolicyRunner  # noqa: I001
 from .distillation_runner import DistillationRunner
-from .on_policy_runner import OnPolicyRunner
 
 __all__ = ["DistillationRunner", "OnPolicyRunner"]