fix comments

Author: ClemensSchwarke

config/example_config.yaml

@@ -1,21 +1,21 @@
 runner:
   class_name: OnPolicyRunner
-  # -- general
-  num_steps_per_env: 24  # number of steps per environment per iteration
-  max_iterations: 1500  # number of policy updates
+  # General
+  num_steps_per_env: 24  # Number of steps per environment per iteration
+  max_iterations: 1500  # Number of policy updates
   seed: 1
-  # -- observations
-  obs_groups: {"policy": ["policy"], "critic": ["policy", "privileged"]} # maps observation groups to sets. See `vec_env.py` for more information
-  # -- logging parameters
-  save_interval: 50  # check for potential saves every `save_interval` iterations
+  # Observations
+  obs_groups: {"policy": ["policy"], "critic": ["policy", "privileged"]} # Maps observation groups to sets. See `vec_env.py` for more information
+  # Logging parameters
+  save_interval: 50  # Check for potential saves every `save_interval` iterations
   experiment_name: walking_experiment
   run_name: ""
-  # -- logging writer
+  # Logging writer
   logger: tensorboard  # tensorboard, neptune, wandb
   neptune_project: legged_gym
   wandb_project: legged_gym
 
-  # -- policy
+  # Policy
   policy:
     class_name: ActorCritic
     activation: elu
@@ -27,45 +27,45 @@ runner:
     noise_std_type: "scalar"  # 'scalar' or 'log'
     state_dependent_std: false
 
-  # -- algorithm
+  # Algorithm
   algorithm:
     class_name: PPO
-    # -- training
+    # Training
     learning_rate: 0.001
     num_learning_epochs: 5
     num_mini_batches: 4  # mini batch size = num_envs * num_steps / num_mini_batches
     schedule: adaptive  # adaptive, fixed
-    # -- value function
+    # Value function
     value_loss_coef: 1.0
     clip_param: 0.2
     use_clipped_value_loss: true
-    # -- surrogate loss
+    # Surrogate loss
     desired_kl: 0.01
     entropy_coef: 0.01
     gamma: 0.99
     lam: 0.95
     max_grad_norm: 1.0
-    # -- miscellaneous
+    # Miscellaneous
     normalize_advantage_per_mini_batch: false
 
-    # -- random network distillation
+    # Random network distillation
     rnd_cfg:
-        weight: 0.0  # initial weight of the RND reward
-        weight_schedule: null # note: this is a dictionary with a required key called "mode". Please check the RND module for more information
-        reward_normalization: false  # whether to normalize RND reward
-        # -- learning parameters
-        learning_rate: 0.001  # learning rate for RND
-        # -- network parameters
-        num_outputs: 1  # number of outputs of RND network. Note: if -1, then the network will use dimensions of the observation
-        predictor_hidden_dims: [-1] # hidden dimensions of predictor network
-        target_hidden_dims: [-1]  # hidden dimensions of target network
+        weight: 0.0  # Initial weight of the RND reward
+        weight_schedule: null # This is a dictionary with a required key called "mode". Please check the RND module for more information
+        reward_normalization: false  # Whether to normalize RND reward
+        # Learning parameters
+        learning_rate: 0.001  # Learning rate for RND
+        # Network parameters
+        num_outputs: 1  # Number of outputs of RND network. Note: if -1, then the network will use dimensions of the observation
+        predictor_hidden_dims: [-1] # Hidden dimensions of predictor network
+        target_hidden_dims: [-1]  # Hidden dimensions of target network
 
-    # -- symmetry augmentation
+    # Symmetry augmentation
     symmetry_cfg:
-      use_data_augmentation: true  # this adds symmetric trajectories to the batch
-      use_mirror_loss: false  # this adds symmetry loss term to the loss function
-      data_augmentation_func: null # string containing the module and function name to import
-      # example: "legged_gym.envs.locomotion.anymal_c.symmetry:get_symmetric_states"
+      use_data_augmentation: true  # This adds symmetric trajectories to the batch
+      use_mirror_loss: false  # This adds symmetry loss term to the loss function
+      data_augmentation_func: null # String containing the module and function name to import
+      # Example: "legged_gym.envs.locomotion.anymal_c.symmetry:get_symmetric_states"
       #
       # .. code-block:: python
       #
@@ -74,4 +74,4 @@ runner:
       #        obs: Optional[torch.Tensor] = None, actions: Optional[torch.Tensor] = None, cfg: "BaseEnvCfg" = None, obs_type: str = "policy"
       #     ) -> Tuple[torch.Tensor, torch.Tensor]:
       #
-      mirror_loss_coeff: 0.0 # coefficient for symmetry loss term. If 0, no symmetry loss is used
+      mirror_loss_coeff: 0.0 # Coefficient for symmetry loss term. If 0, no symmetry loss is used

rsl_rl/algorithms/distillation.py

@@ -31,9 +31,10 @@ def __init__(
         # Distributed training parameters
         multi_gpu_cfg: dict | None = None,
     ) -> None:
-        # device-related parameters
+        # Device-related parameters
         self.device = device
         self.is_multi_gpu = multi_gpu_cfg is not None
+
         # Multi-GPU parameters
         if multi_gpu_cfg is not None:
             self.gpu_global_rank = multi_gpu_cfg["global_rank"]
@@ -42,25 +43,25 @@ def __init__(
             self.gpu_global_rank = 0
             self.gpu_world_size = 1
 
-        # distillation components
+        # Distillation components
         self.policy = policy
         self.policy.to(self.device)
-        self.storage = None  # initialized later
+        self.storage = None  # Initialized later
 
-        # initialize the optimizer
+        # Initialize the optimizer
         self.optimizer = resolve_optimizer(optimizer)(self.policy.parameters(), lr=learning_rate)
 
-        # initialize the transition
+        # Initialize the transition
         self.transition = RolloutStorage.Transition()
         self.last_hidden_states = None
 
-        # distillation parameters
+        # Distillation parameters
         self.num_learning_epochs = num_learning_epochs
         self.gradient_length = gradient_length
         self.learning_rate = learning_rate
         self.max_grad_norm = max_grad_norm
 
-        # initialize the loss function
+        # Initialize the loss function
         loss_fn_dict = {
             "mse": nn.functional.mse_loss,
             "huber": nn.functional.huber_loss,
@@ -80,7 +81,7 @@ def init_storage(
         obs: TensorDict,
         actions_shape: tuple[int],
     ) -> None:
-        # create rollout storage
+        # Create rollout storage
         self.storage = RolloutStorage(
             training_type,
             num_envs,
@@ -91,23 +92,23 @@ def init_storage(
         )
 
     def act(self, obs: TensorDict) -> torch.Tensor:
-        # compute the actions
+        # Compute the actions
         self.transition.actions = self.policy.act(obs).detach()
         self.transition.privileged_actions = self.policy.evaluate(obs).detach()
-        # record the observations
+        # Record the observations
         self.transition.observations = obs
         return self.transition.actions
 
     def process_env_step(
         self, obs: TensorDict, rewards: torch.Tensor, dones: torch.Tensor, extras: dict[str, torch.Tensor]
     ) -> None:
-        # update the normalizers
+        # Update the normalizers
         self.policy.update_normalization(obs)
 
-        # record the rewards and dones
+        # Record the rewards and dones
         self.transition.rewards = rewards
         self.transition.dones = dones
-        # record the transition
+        # Record the transition
         self.storage.add_transitions(self.transition)
         self.transition.clear()
         self.policy.reset(dones)
@@ -122,18 +123,18 @@ def update(self) -> dict[str, float]:
             self.policy.reset(hidden_states=self.last_hidden_states)
             self.policy.detach_hidden_states()
             for obs, _, privileged_actions, dones in self.storage.generator():
-                # inference the student for gradient computation
+                # Inference of the student for gradient computation
                 actions = self.policy.act_inference(obs)
 
-                # behavior cloning loss
+                # Behavior cloning loss
                 behavior_loss = self.loss_fn(actions, privileged_actions)
 
-                # total loss
+                # Total loss
                 loss = loss + behavior_loss
                 mean_behavior_loss += behavior_loss.item()
                 cnt += 1
 
-                # gradient step
+                # Gradient step
                 if cnt % self.gradient_length == 0:
                     self.optimizer.zero_grad()
                     loss.backward()
@@ -145,7 +146,7 @@ def update(self) -> dict[str, float]:
                     self.policy.detach_hidden_states()
                     loss = 0
 
-                # reset dones
+                # Reset dones
                 self.policy.reset(dones.view(-1))
                 self.policy.detach_hidden_states(dones.view(-1))
 
@@ -154,22 +155,18 @@ def update(self) -> dict[str, float]:
         self.last_hidden_states = self.policy.get_hidden_states()
         self.policy.detach_hidden_states()
 
-        # construct the loss dictionary
+        # Construct the loss dictionary
         loss_dict = {"behavior": mean_behavior_loss}
 
         return loss_dict
 
-    """
-    Helper functions
-    """
-
     def broadcast_parameters(self) -> None:
         """Broadcast model parameters to all GPUs."""
-        # obtain the model parameters on current GPU
+        # Obtain the model parameters on current GPU
         model_params = [self.policy.state_dict()]
-        # broadcast the model parameters
+        # Broadcast the model parameters
         torch.distributed.broadcast_object_list(model_params, src=0)
-        # load the model parameters on all GPUs from source GPU
+        # Load the model parameters on all GPUs from source GPU
         self.policy.load_state_dict(model_params[0])
 
     def reduce_parameters(self) -> None:
@@ -188,7 +185,7 @@ def reduce_parameters(self) -> None:
         for param in self.policy.parameters():
             if param.grad is not None:
                 numel = param.numel()
-                # copy data back from shared buffer
+                # Copy data back from shared buffer
                 param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data))
-                # update the offset for the next parameter
+                # Update the offset for the next parameter
                 offset += numel