[Question] Adding a GRU in Dict_Discrete_PPO Issaclab

[shiven2001]

hello, I am trying to do this to add a GRU and I am running in an issue. I followed the Toni-SM/skrl#276 thread but I get this error:

  Traceback (most recent call last):
    File "/home/shiven/IsaacLab/source/isaaclab_tasks/isaaclab_tasks/direct/aod_float/skrl_discrete_train.py", line 228, in <module>
      trainer.train()
    File "/home/shiven/env_isaaclab/lib/python3.10/site-packages/skrl/trainers/torch/sequential.py", line 86, in train
      self.single_agent_train()
    File "/home/shiven/env_isaaclab/lib/python3.10/site-packages/skrl/trainers/torch/base.py", line 193, in single_agent_train
      actions = self.agents.act(states, timestep=timestep, timesteps=self.timesteps)[0]
    File "/home/shiven/env_isaaclab/lib/python3.10/site-packages/skrl/agents/torch/ppo/ppo.py", line 243, in act
      actions, log_prob, outputs = self.policy.act({"states": self._state_preprocessor(states)}, role="policy")
    File "/home/shiven/IsaacLab/source/isaaclab_tasks/isaaclab_tasks/direct/aod_float/skrl_discrete_train.py", line 96, in act
      return CategoricalMixin.act(self, inputs, role)
    File "/home/shiven/env_isaaclab/lib/python3.10/site-packages/skrl/models/torch/categorical.py", line 89, in act
      net_output, outputs = self.compute(inputs, role)
    File "/home/shiven/IsaacLab/source/isaaclab_tasks/isaaclab_tasks/direct/aod_float/skrl_discrete_train.py", line 103, in compute
      hidden_states = inputs["rnn"][0]
  KeyError: 'rnn'
  [skrl:INFO] Closing environment
  [skrl:INFO] Environment closed

I am using a dict as observation and discrete action space. What is the issue here? Also another question, it seems that GRU is reset on termination. Will it also reset on truncation? I am using this as my get_dones in DirectRL:

    def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]:
            time_out = self.episode_length_buf >= self.max_episode_length - 1
            # No detection
            termination = self.out_of_bounds
            print(termination, flush=True)
    
            return termination, time_out  # termination, truncation

There is termination and truncation conditions. I don't know if the GRU is reset on the time_out as well or do I need to add some code for that? Thanks

Here is the skrl code:

  import torch
  import torch.nn as nn
  
  # import the skrl components to build the RL system
  from skrl.agents.torch.ppo import PPO, PPO_DEFAULT_CONFIG
  from skrl.envs.loaders.torch import load_isaaclab_env
  from skrl.envs.wrappers.torch import wrap_env
  from skrl.memories.torch import RandomMemory
  from skrl.models.torch import DeterministicMixin, CategoricalMixin, Model
  from skrl.resources.preprocessors.torch import RunningStandardScaler
  from skrl.resources.schedulers.torch import KLAdaptiveRL
  from skrl.trainers.torch import SequentialTrainer
  from skrl.utils import set_seed
  

  # Helper function to extract dimensions from gym spaces
  import gym
  
  
  def get_space_dim(space):
      # Handle Dict spaces by summing the dimensions of all subspaces
      if hasattr(space, "spaces") and isinstance(space.spaces, dict):
          return sum(get_space_dim(subspace) for subspace in space.spaces.values())
      elif hasattr(space, "shape") and space.shape is not None and len(space.shape) > 0:
          return int(torch.prod(torch.tensor(space.shape)))
      elif hasattr(space, "n"):
          return int(space.n)
      else:
          raise ValueError(f"Unknown space type: {space}")
  
  
  # seed for reproducibility
  set_seed(42)  # e.g. `set_seed(42)` for fixed seed
  
  
  # define shared model (stochastic and deterministic models) using mixins
  class Shared(CategoricalMixin, DeterministicMixin, Model):
      def __init__(
          self,
          observation_space,
          action_space,
          device,
          unnormalized_log_prob=True,
          clip_actions=False,
          sequence_length=128,
      ):
          Model.__init__(self, observation_space, action_space, device)
          CategoricalMixin.__init__(self, unnormalized_log_prob)
          DeterministicMixin.__init__(self, clip_actions)
  
          # Extract dimensions
          self.num_observations = get_space_dim(observation_space)
          self.num_actions = get_space_dim(action_space)
          self.sequence_length = sequence_length
          self.num_layers = 1
  
          self.net = nn.Sequential(
              nn.Linear(self.num_observations, 256),
              nn.ELU(),
              nn.Linear(256, 128),
              nn.ELU(),
              nn.Linear(128, 64),
              nn.ELU(),
          )
  
          self.gru = nn.GRU(
              input_size=self.num_observations,
              hidden_size=64,
              num_layers=self.num_layers,
              batch_first=True,
          )  # batch_first -> (batch, sequence, features)
  
          self.mean_layer = nn.Linear(64, self.num_actions)
          self.value_layer = nn.Linear(64, 1)
  
          # For RNN specification
          self.hidden_size = 64
          self.num_layers = 1
  
      def get_specification(self):
          # batch size (N) is the number of envs
          return {
              "rnn": {
                  "sequence_length": self.sequence_length,
                  "sizes": [(self.num_layers, self.num_envs, self.hidden_size)],
              }
          }  # hidden states (D ∗ num_layers, N, Hout)
  
      def act(self, inputs, role):
          if role == "policy":
              return CategoricalMixin.act(self, inputs, role)
          elif role == "value":
              return DeterministicMixin.act(self, inputs, role)
  
      def compute(self, inputs, role):
          states = inputs["states"]
          terminated = inputs.get("terminated", None)
          hidden_states = inputs["rnn"][0]
  
          rnn_output, rnn_dict = self.rnn_rollout(states, terminated, hidden_states)
          if role == "policy":
              return self.mean_layer(rnn_output), rnn_dict
          elif role == "value":
              return self.value_layer(rnn_output), rnn_dict
  
      def rnn_rollout(self, states, terminated, hidden_states):
          if self.training:
              rnn_input = states.view(
                  -1, self.sequence_length, states.shape[-1]
              )  # (N, L, Hin): N=batch_size, L=sequence_length
              hidden_states = hidden_states.view(
                  self.num_layers, -1, self.sequence_length, hidden_states.shape[-1]
              )  # (D * num_layers, N, L, Hout)
              # get the hidden states corresponding to the initial sequence
              hidden_states = hidden_states[
                  :, :, 0, :
              ].contiguous()  # (D * num_layers, N, Hout)
  
              # reset the RNN state in the middle of a sequence
              if terminated is not None and torch.any(terminated):
                  rnn_outputs = []
                  terminated = terminated.view(-1, self.sequence_length)
                  indexes = (
                      [0]
                      + (
                          terminated[:, :-1].any(dim=0).nonzero(as_tuple=True)[0] + 1
                      ).tolist()
                      + [self.sequence_length]
                  )
  
                  for i in range(len(indexes) - 1):
                      i0, i1 = indexes[i], indexes[i + 1]
                      rnn_output, hidden_states = self.gru(
                          rnn_input[:, i0:i1, :], hidden_states
                      )
                      hidden_states[:, (terminated[:, i1 - 1]), :] = 0
                      rnn_outputs.append(rnn_output)
  
                  rnn_output = torch.cat(rnn_outputs, dim=1)
              # no need to reset the RNN state in the sequence
              else:
                  rnn_output, hidden_states = self.gru(rnn_input, hidden_states)
          # rollout
          else:
              rnn_input = states.view(
                  -1, 1, states.shape[-1]
              )  # (N, L, Hin): N=num_envs, L=1
              rnn_output, hidden_states = self.gru(rnn_input, hidden_states)
  
          # flatten the RNN output
          rnn_output = torch.flatten(
              rnn_output, start_dim=0, end_dim=1
          )  # (N, L, D ∗ Hout) -> (N * L, D ∗ Hout)
  
          # action_space
          return self.net(rnn_output), {"rnn": [hidden_states]}
  
  
  # load and wrap the Isaac Lab environment
  env = load_isaaclab_env(task_name="AOD-Float-Discrete-v0")
  env = wrap_env(env)
  
  device = env.device
  
  # instantiate a memory as rollout buffer (any memory can be used for this)
  memory = RandomMemory(memory_size=16, num_envs=env.num_envs, device=device)
  
  # instantiate the agent's models (function approximators).
  # PPO requires 2 models, visit its documentation for more details
  # https://skrl.readthedocs.io/en/latest/api/agents/ppo.html#models
  models = {}
  models["policy"] = Shared(env.observation_space, env.action_space, device)
  models["value"] = models["policy"]  # same instance: shared model
  
  # configure and instantiate the agent (visit its documentation to see all the options)
  # https://skrl.readthedocs.io/en/latest/api/agents/ppo.html#configuration-and-hyperparameters
  cfg = PPO_DEFAULT_CONFIG.copy()
  cfg["rollouts"] = 16  # memory_size
  cfg["learning_epochs"] = 8
  cfg["mini_batches"] = 4  # 16 * 1024 / 4096
  cfg["discount_factor"] = 0.99
  cfg["lambda"] = 0.95
  cfg["learning_rate"] = 3e-4
  cfg["learning_rate_scheduler"] = KLAdaptiveRL
  cfg["learning_rate_scheduler_kwargs"] = {"kl_threshold": 0.008}
  cfg["random_timesteps"] = 0
  cfg["learning_starts"] = 0
  cfg["grad_norm_clip"] = 1.0
  cfg["ratio_clip"] = 0.2
  cfg["value_clip"] = 0.2
  cfg["clip_predicted_values"] = True
  cfg["entropy_loss_scale"] = 0.0
  cfg["value_loss_scale"] = 1.0
  cfg["kl_threshold"] = 0
  cfg["rewards_shaper"] = lambda rewards, *args, **kwargs: rewards * 0.1
  cfg["time_limit_bootstrap"] = True
  cfg["state_preprocessor"] = RunningStandardScaler
  cfg["state_preprocessor_kwargs"] = {
      "size": get_space_dim(env.observation_space),
      "device": device,
  }
  cfg["value_preprocessor"] = RunningStandardScaler
  cfg["value_preprocessor_kwargs"] = {"size": 1, "device": device}
  # logging to TensorBoard and write checkpoints (in timesteps)
  cfg["experiment"]["write_interval"] = 40
  cfg["experiment"]["checkpoint_interval"] = 400
  cfg["experiment"]["directory"] = "runs/torch/AOD-Float-Discrete-v0"
  
  agent = PPO(
      models=models,
      memory=memory,
      cfg=cfg,
      observation_space=env.observation_space,
      action_space=env.action_space,
      device=device,
  )
  
  # configure and instantiate the RL trainer
  cfg_trainer = {"timesteps": 20000, "headless": True}
  trainer = SequentialTrainer(cfg=cfg_trainer, env=env, agents=agent)
  
  # start training
  trainer.train()

[Toni-SM]

Hi @shiven2001

As mentioned in the this skrl's docs note, support for recurrent neural networks (RNN, LSTM, GRU and any other variant) is implemented in a separate file.

In this case your import should be as follow:

from skrl.agents.torch.ppo import PPO_DEFAULT_CONFIG
from skrl.agents.torch.ppo import PPO_RNN as PPO

[shiven2001]

Hi @shiven2001

As mentioned in the this skrl's docs note, support for recurrent neural networks (RNN, LSTM, GRU and any other variant) is implemented in a separate file.

In this case your import should be as follow:

from skrl.agents.torch.ppo import PPO_DEFAULT_CONFIG
from skrl.agents.torch.ppo import PPO_RNN as PPO

Sorry I didn't see the PPO_RNN as PPO part. It runs properly now. I am running it using a CategoricalMixin as my action space is discrete. I will provide the code as a reference in case anyone want to set up an discrete action space GRU PPO policy. I have two more questions:

1. My observation space is a dict. In the example conf file, it had:

        network:
              - name: features_extractor
                input: permute(OBSERVATIONS["camera"], (0, 3, 1, 2))  # PyTorch NHWC -> NCHW. Warning: don't permute for JAX since it expects NHWC
                layers:
                  - conv2d: {out_channels: 32, kernel_size: 8, stride: 4, padding: 0}
                  - conv2d: {out_channels: 64, kernel_size: 4, stride: 2, padding: 0}
                  - conv2d: {out_channels: 64, kernel_size: 3, stride: 1, padding: 0}
                  - flatten
                  - linear: 512
                  - linear: 16
                activations: [relu, relu, relu, null, tanh, tanh]
              - name: net
                input: concatenate([features_extractor, OBSERVATIONS["joint-velocities"]])
                layers: [32, 32]
                activations: elu
            output: ACTIONS
   

It seems I can use states = inputs["states"] directly in INPUTS instead of the concatenate([features_extractor, OBSERVATIONS["joint-velocities"]]) from the environment. But I want to add an image input as use the

      - name: features_extractor
                    input: permute(OBSERVATIONS["camera"], (0, 3, 1, 2))  # PyTorch NHWC -> NCHW. Warning: don't permute for JAX since it expects NHWC
                    layers:
                      - conv2d: {out_channels: 32, kernel_size: 8, stride: 4, padding: 0}
                      - conv2d: {out_channels: 64, kernel_size: 4, stride: 2, padding: 0}
                      - conv2d: {out_channels: 64, kernel_size: 3, stride: 1, padding: 0}
                      - flatten
                      - linear: 512
                      - linear: 16
                    activations: [relu, relu, relu, null, tanh, tanh]

on it. I am unsure how the states = inputs["states"] work with isaaclab. If I use states = inputs["states"], and if I have the image in my observation space, will it also be included in the states = inputs["states"]? If yes, how do I select the image first and run a features_extractor on it before and then adding the extracted features and the rest of the observations together in GRU instead of the image and rest of the observations?

2. About the termination question I had. it seems that GRU is reset on termination. Will it also reset on truncation? I am using this as my get_dones in DirectRL:

        def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]:
                time_out = self.episode_length_buf >= self.max_episode_length - 1
                # No detection
                termination = self.out_of_bounds
                print(termination, flush=True)
        
                return termination, time_out  # termination, truncation
   

There is termination and truncation conditions. I don't know if the GRU is reset on the time_out as well or do I need to add some code for that? Thanks

FULL CODE:

    import torch
    import torch.nn as nn
    
    # import the skrl components to build the RL system
    from skrl.agents.torch.ppo import PPO_DEFAULT_CONFIG
    from skrl.agents.torch.ppo import PPO_RNN as PPO
    from skrl.envs.loaders.torch import load_isaaclab_env
    from skrl.envs.wrappers.torch import wrap_env
    from skrl.memories.torch import RandomMemory
    from skrl.models.torch import DeterministicMixin, CategoricalMixin, Model
    from skrl.resources.preprocessors.torch import RunningStandardScaler
    from skrl.resources.schedulers.torch import KLAdaptiveRL
    from skrl.trainers.torch import SequentialTrainer
    from skrl.utils import set_seed
    
    import numpy as np
    import torch
    import torch.nn as nn
    import gymnasium as gym
    
    # import the skrl components to build the RL system
    from skrl.agents.torch.ppo import PPO_DEFAULT_CONFIG
    from skrl.agents.torch.ppo import PPO_RNN as PPO
    from skrl.envs.wrappers.torch import wrap_env
    from skrl.memories.torch import RandomMemory
    from skrl.models.torch import DeterministicMixin, CategoricalMixin, Model
    from skrl.resources.preprocessors.torch import RunningStandardScaler
    from skrl.resources.schedulers.torch import KLAdaptiveRL
    from skrl.trainers.torch import SequentialTrainer
    from skrl.utils import set_seed
    
    
    # seed for reproducibility
    set_seed(42)  # e.g. `set_seed(42)` for fixed seed
    
    
    # define models (stochastic and deterministic models) using mixins
    class Policy(CategoricalMixin, Model):
        def __init__(
            self,
            observation_space,
            action_space,
            device,
            unnormalized_log_prob=True,
            num_envs=1,
            num_layers=1,
            hidden_size=64,
            sequence_length=128,
        ):
            Model.__init__(self, observation_space, action_space, device)
            CategoricalMixin.__init__(self, unnormalized_log_prob)
    
            self.num_envs = num_envs
            self.num_layers = num_layers
            self.hidden_size = hidden_size  # Hout
            self.sequence_length = sequence_length
    
            self.gru = nn.GRU(
                input_size=self.num_observations,
                hidden_size=self.hidden_size,
                num_layers=self.num_layers,
                batch_first=True,
            )  # batch_first -> (batch, sequence, features)
    
            self.net = nn.Sequential(
                nn.Linear(self.hidden_size, 256),
                nn.ReLU(),
                nn.Linear(256, 128),
                nn.ReLU(),
                nn.Linear(128, 64),
                nn.ReLU(),
                nn.Linear(64, self.num_actions),
            )
    
        def get_specification(self):
            # batch size (N) is the number of envs
            return {
                "rnn": {
                    "sequence_length": self.sequence_length,
                    "sizes": [(self.num_layers, self.num_envs, self.hidden_size)],
                }
            }  # hidden states (D ∗ num_layers, N, Hout)
    
        def compute(self, inputs, role):
            states = inputs["states"]
            terminated = inputs.get("terminated", None)
            hidden_states = inputs["rnn"][0]
    
            # training
            if self.training:
                rnn_input = states.view(
                    -1, self.sequence_length, states.shape[-1]
                )  # (N, L, Hin): N=batch_size, L=sequence_length
                hidden_states = hidden_states.view(
                    self.num_layers, -1, self.sequence_length, hidden_states.shape[-1]
                )  # (D * num_layers, N, L, Hout)
                # get the hidden states corresponding to the initial sequence
                hidden_states = hidden_states[
                    :, :, 0, :
                ].contiguous()  # (D * num_layers, N, Hout)
    
                # reset the RNN state in the middle of a sequence
                if terminated is not None and torch.any(terminated):
                    rnn_outputs = []
                    terminated = terminated.view(-1, self.sequence_length)
                    indexes = (
                        [0]
                        + (
                            terminated[:, :-1].any(dim=0).nonzero(as_tuple=True)[0] + 1
                        ).tolist()
                        + [self.sequence_length]
                    )
    
                    for i in range(len(indexes) - 1):
                        i0, i1 = indexes[i], indexes[i + 1]
                        rnn_output, hidden_states = self.gru(
                            rnn_input[:, i0:i1, :], hidden_states
                        )
                        hidden_states[:, (terminated[:, i1 - 1]), :] = 0
                        rnn_outputs.append(rnn_output)
    
                    rnn_output = torch.cat(rnn_outputs, dim=1)
                # no need to reset the RNN state in the sequence
                else:
                    rnn_output, hidden_states = self.gru(rnn_input, hidden_states)
            # rollout
            else:
                rnn_input = states.view(
                    -1, 1, states.shape[-1]
                )  # (N, L, Hin): N=num_envs, L=1
                rnn_output, hidden_states = self.gru(rnn_input, hidden_states)
    
            # flatten the RNN output
            rnn_output = torch.flatten(
                rnn_output, start_dim=0, end_dim=1
            )  # (N, L, D ∗ Hout) -> (N * L, D ∗ Hout)
            # Number of environments (or batch size during training).
            # Sequence length (1 during rollout, sequence_length during training).
            # GRU's output dimension, where Hout is self.hidden_size.
            # After flattening, rnn_output has shape (N * L, hidden_size).
    
            return self.net(rnn_output), {"rnn": [hidden_states]}
    
    
    class Value(DeterministicMixin, Model):
        def __init__(
            self,
            observation_space,
            action_space,
            device,
            clip_actions=False,
            num_envs=1,
            num_layers=1,
            hidden_size=64,
            sequence_length=128,
        ):
            Model.__init__(self, observation_space, action_space, device)
            DeterministicMixin.__init__(self, clip_actions)
    
            self.num_envs = num_envs
            self.num_layers = num_layers
            self.hidden_size = hidden_size  # Hout
            self.sequence_length = sequence_length
    
            self.gru = nn.GRU(
                input_size=self.num_observations,
                hidden_size=self.hidden_size,
                num_layers=self.num_layers,
                batch_first=True,
            )  # batch_first -> (batch, sequence, features)
    
            self.net = nn.Sequential(
                nn.Linear(self.hidden_size, 256),
                nn.ReLU(),
                nn.Linear(256, 128),
                nn.ReLU(),
                nn.Linear(128, 64),
                nn.ReLU(),
                nn.Linear(64, 1),
            )
    
        def get_specification(self):
            # batch size (N) is the number of envs
            return {
                "rnn": {
                    "sequence_length": self.sequence_length,
                    "sizes": [(self.num_layers, self.num_envs, self.hidden_size)],
                }
            }  # hidden states (D ∗ num_layers, N, Hout)
    
        def compute(self, inputs, role):
            states = inputs["states"]
            terminated = inputs.get("terminated", None)
            hidden_states = inputs["rnn"][0]
    
            # training
            if self.training:
                rnn_input = states.view(
                    -1, self.sequence_length, states.shape[-1]
                )  # (N, L, Hin): N=batch_size, L=sequence_length
    
                hidden_states = hidden_states.view(
                    self.num_layers, -1, self.sequence_length, hidden_states.shape[-1]
                )  # (D * num_layers, N, L, Hout)
                # get the hidden states corresponding to the initial sequence
                hidden_states = hidden_states[
                    :, :, 0, :
                ].contiguous()  # (D * num_layers, N, Hout)
    
                # reset the RNN state in the middle of a sequence
                if terminated is not None and torch.any(terminated):
                    rnn_outputs = []
                    terminated = terminated.view(-1, self.sequence_length)
                    indexes = (
                        [0]
                        + (
                            terminated[:, :-1].any(dim=0).nonzero(as_tuple=True)[0] + 1
                        ).tolist()
                        + [self.sequence_length]
                    )
    
                    for i in range(len(indexes) - 1):
                        i0, i1 = indexes[i], indexes[i + 1]
                        rnn_output, hidden_states = self.gru(
                            rnn_input[:, i0:i1, :], hidden_states
                        )
                        hidden_states[:, (terminated[:, i1 - 1]), :] = 0
                        rnn_outputs.append(rnn_output)
    
                    rnn_output = torch.cat(rnn_outputs, dim=1)
                # no need to reset the RNN state in the sequence
                else:
                    rnn_output, hidden_states = self.gru(rnn_input, hidden_states)
            # rollout
            else:
                rnn_input = states.view(
                    -1, 1, states.shape[-1]
                )  # (N, L, Hin): N=num_envs, L=1
                rnn_output, hidden_states = self.gru(rnn_input, hidden_states)
    
            # flatten the RNN output
            rnn_output = torch.flatten(
                rnn_output, start_dim=0, end_dim=1
            )  # (N, L, D ∗ Hout) -> (N * L, D ∗ Hout)
    
            return self.net(rnn_output), {"rnn": [hidden_states]}
    
    
    # load and wrap the Isaac Lab environment
    env = load_isaaclab_env(task_name="AOD-Float-Discrete-v0")
    env = wrap_env(env)
    
    device = env.device
    
    # instantiate a memory as rollout buffer (any memory can be used for this)
    memory = RandomMemory(memory_size=1024, num_envs=env.num_envs, device=device)
    
    # instantiate the agent's models (function approximators).
    # PPO requires 2 models, visit its documentation for more details
    # https://skrl.readthedocs.io/en/latest/api/agents/ppo.html#models
    
    models = {}
    models["policy"] = Policy(
        observation_space=env.observation_space,
        action_space=env.action_space,
        device=device,
        num_envs=env.num_envs,
    )
    models["value"] = Value(
        observation_space=env.observation_space,
        action_space=env.action_space,
        device=device,
        num_envs=env.num_envs,
    )
    
    # configure and instantiate the agent (visit its documentation to see all the options)
    # https://skrl.readthedocs.io/en/latest/api/agents/ppo.html#configuration-and-hyperparameters
    cfg = PPO_DEFAULT_CONFIG.copy()
    cfg["rollouts"] = 1024  # memory_size
    cfg["learning_epochs"] = 10
    cfg["mini_batches"] = 4  # 16 * 1024 / 4096
    cfg["discount_factor"] = 0.99
    cfg["lambda"] = 0.95
    cfg["learning_rate"] = 1e-3
    cfg["learning_rate_scheduler"] = KLAdaptiveRL
    cfg["learning_rate_scheduler_kwargs"] = {"kl_threshold": 0.008}
    cfg["random_timesteps"] = 0
    cfg["learning_starts"] = 0
    cfg["grad_norm_clip"] = 0.5
    cfg["ratio_clip"] = 0.2
    cfg["value_clip"] = 0.2
    cfg["clip_predicted_values"] = True
    cfg["entropy_loss_scale"] = 0.0
    cfg["value_loss_scale"] = 1.0
    cfg["kl_threshold"] = 0
    cfg["rewards_shaper"] = lambda rewards, *args, **kwargs: rewards * 0.1
    cfg["time_limit_bootstrap"] = True
    cfg["state_preprocessor"] = RunningStandardScaler
    cfg["state_preprocessor_kwargs"] = {"size": env.observation_space, "device": device}
    cfg["value_preprocessor"] = RunningStandardScaler
    cfg["value_preprocessor_kwargs"] = {"size": 1, "device": device}
    # logging to TensorBoard and write checkpoints (in timesteps)
    cfg["experiment"]["write_interval"] = 50
    cfg["experiment"]["checkpoint_interval"] = 500
    cfg["experiment"]["directory"] = "runs/torch/AOD-Float-Discrete-v0-GRU"
    
    agent = PPO(
        models=models,
        memory=memory,
        cfg=cfg,
        observation_space=env.observation_space,
        action_space=env.action_space,
        device=device,
    )
    
    # configure and instantiate the RL trainer
    cfg_trainer = {"timesteps": 20000, "headless": True}
    trainer = SequentialTrainer(cfg=cfg_trainer, env=env, agents=agent)
    
    # start training
    trainer.train()
    
    # # ---------------------------------------------------------
    # # comment the code above: `trainer.train()`, and...
    # # uncomment the following lines to evaluate a trained agent
    # # ---------------------------------------------------------
    # from skrl.utils.huggingface import download_model_from_huggingface
    
    # # download the trained agent's checkpoint from Hugging Face Hub and load it
    # path = download_model_from_huggingface("skrl/IsaacOrbit-Isaac-Ant-v0-PPO", filename="agent.pt")
    # agent.load(path)
    
    # # start evaluation
    # trainer.eval()

[RandomOakForest]

Thank you for following up, I'll move this post to our Discussions for the team and others to reply.