import gymnasium as gym import numpy as np 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.wrappers.torch import wrap_env from skrl.memories.torch import RandomMemory from skrl.models.torch import DeterministicMixin, GaussianMixin, 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() # e.g. `set_seed(42)` for fixed seed # define models (stochastic and deterministic models) using mixins class Policy(GaussianMixin, Model): def __init__(self, observation_space, action_space, device, clip_actions=False, clip_log_std=True, min_log_std=-20, max_log_std=2, reduction="sum", num_envs=1, num_layers=1, hidden_size=64, sequence_length=128): Model.__init__(self, observation_space, action_space, device) GaussianMixin.__init__(self, clip_actions, clip_log_std, min_log_std, max_log_std, reduction) self.num_envs = num_envs self.num_layers = num_layers self.hidden_size = hidden_size # Hout self.sequence_length = sequence_length self.rnn = nn.RNN(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, 64), nn.ReLU(), nn.Linear(64, self.num_actions)) self.log_std_parameter = nn.Parameter(torch.zeros(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.rnn(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.rnn(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.rnn(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) # Pendulum-v1 action_space is -2 to 2 return 2 * torch.tanh(self.net(rnn_output)), self.log_std_parameter, {"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.rnn = nn.RNN(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, 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.rnn(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.rnn(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.rnn(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]} # environment observation wrapper used to mask velocity. Adapted from rl_zoo3 (rl_zoo3/wrappers.py) class NoVelocityWrapper(gym.ObservationWrapper): def observation(self, observation): # observation: x, y, angular velocity return observation * np.array([1, 1, 0]) gym.register(id="PendulumNoVel-v1", entry_point=lambda: NoVelocityWrapper(gym.make("Pendulum-v1"))) # load and wrap the gymnasium environment env = gym.make_vec("PendulumNoVel-v1", num_envs=4, vectorization_mode="sync") 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(env.observation_space, env.action_space, device, clip_actions=True, num_envs=env.num_envs) models["value"] = Value(env.observation_space, env.action_space, 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"] = 32 cfg["discount_factor"] = 0.9 cfg["lambda"] = 0.95 cfg["learning_rate"] = 1e-3 cfg["learning_rate_scheduler"] = KLAdaptiveRL cfg["learning_rate_scheduler_kwargs"] = {"kl_threshold": 0.008} cfg["grad_norm_clip"] = 0.5 cfg["ratio_clip"] = 0.2 cfg["value_clip"] = 0.2 cfg["clip_predicted_values"] = False cfg["entropy_loss_scale"] = 0.0 cfg["value_loss_scale"] = 0.5 cfg["kl_threshold"] = 0 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"] = 500 cfg["experiment"]["checkpoint_interval"] = 5000 cfg["experiment"]["directory"] = "runs/torch/PendulumNoVel" 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": 100000, "headless": True} trainer = SequentialTrainer(cfg=cfg_trainer, env=env, agents=[agent]) # start training trainer.train()