Add sampling strategies and longer experience buffer

Author: SimBe195

configs/agent/learning_rate/linear_decay.yaml

@@ -1,4 +1,4 @@
 _target_: config.learning_rate.LinearDecayLRConfig
 
-learning_rate: 0.000723
+learning_rate: 0.0007
 decay_iters: 1500

configs/agent/model/actor_critic.yaml

@@ -1,8 +1,6 @@
 _target_: config.model.ActorCriticConfig
 
 action_embedding_size: 128
-num_layers: 4
-num_filters: 128
-kernel_size: 4
-stride: 2
-fc_size: 260
+num_filters: [128, 64, 64, 32]
+kernel_sizes: [9, 5, 5, 3]
+fc_size: 256

configs/agent/ppo.yaml

@@ -1,17 +1,19 @@
 defaults:
   - model: actor_critic
   - learning_rate: linear_decay
+  - sampling_strategy: dist_sample
   - _self_
 
 _target_: config.agent.PPOAgentConfig
 
-gamma: 0.903
-tau: 0.854
+gamma: 0.9
+tau: 0.85
+exp_buffer_size: 512
 epochs_per_update: 5
 total_updates: 0
-batch_size: 256
+batch_size: 512
 clip_param: 0.2
 clip_value: False
-critic_loss_weight: 0.818
-max_entropy_loss_weight: 0.00118
+critic_loss_weight: 0.8
+max_entropy_loss_weight: 0.001
 grad_clip_norm: 0.5

configs/agent/sampling_strategy/adapt_eps_greedy.yaml

@@ -0,0 +1,5 @@
+_target_: config.sampling_strategy.AdaptiveEpsilonGreedySamplingStrategy
+
+initial_epsilon: 0.5
+min_epsilon: 0.01
+decay_factor: 0.99

configs/agent/sampling_strategy/dist_sample.yaml

@@ -0,0 +1 @@
+_target_: config.sampling_strategy.DistSamplingStrategy

configs/agent/sampling_strategy/greedy.yaml

@@ -0,0 +1 @@
+_target_: config.sampling_strategy.GreedySamplingStrategy

configs/environment/base.yaml

@@ -3,7 +3,7 @@ _target_: config.environment.EnvironmentConfig
 env_name: SuperMarioBros-v0
 complex_movement: True
 num_repeat_frames: 4
-num_stack_frames: 3
+num_stack_frames: 2
 clip_top: 0
 clip_bot: 0
 clip_left: 0

configs/neurio_config.yaml

@@ -6,11 +6,11 @@ defaults:
   - _self_
 
 _target_: config.main_config.NeurioConfig
-level: "1-3"
+level: "1-1"
 num_workers: 32
-num_iters: 3000
+num_iters: 500
 steps_per_iter: 128
-save_frequency: 100
+save_frequency: 50
 render: False
 
 agent:
@@ -19,23 +19,21 @@ agent:
   learning_rate:
     decay_iters: ${num_iters}
 
-hydra:
-  sweeper:
-    # ToDo: Add pruner as soon as it's available in the optuna sweeper
-    sampler:
-      n_startup_trials: 10
-    direction: maximize
-    study_name: Neurio-lev-${level}
-    storage: sqlite:///optuna_studies/${hydra.sweeper.study_name}.db
-    n_trials: 1
-    n_jobs: 1
-
-    params:
-      agent.learning_rate.learning_rate: tag(log, interval(0.0002, 0.0009))
-      agent.model.num_filters: choice(128, 256)
-      # agent.model.kernel_size: range(3, 4)
-      # agent.model.fc_size: range(100, 700)
-      # agent.critic_loss_weight: interval(0.5, 1.5)
-      agent.max_entropy_loss_weight: tag(log, interval(0.0001, 0.01))
-      agent.gamma: interval(0.8, 1.0)
-      agent.tau: interval(0.8, 1.0)
+#hydra:
+#  sweeper:
+#    # ToDo: Add pruner as soon as it's available in the optuna sweeper
+#    sampler:
+#      n_startup_trials: 10
+#    direction: maximize
+#    study_name: Neurio-lev-${level}
+#    storage: sqlite:///optuna_studies/${hydra.sweeper.study_name}.db
+#    n_trials: 1
+#    n_jobs: 1
+#
+#    params:
+#      agent.learning_rate.learning_rate: tag(log, interval(0.0002, 0.0009))
+#      agent.exp_buffer_size: choice(128, 256)
+#      agent.batch_size: choice(64, 128, 256, 512)
+#      agent.max_entropy_loss_weight: tag(log, interval(0.0001, 0.01))
+#      agent.gamma: interval(0.8, 1.0)
+#      agent.tau: interval(0.8, 1.0)

src/agents/experience_buffer.py

@@ -1,93 +1,87 @@
-from typing import List
+from collections import deque
+
+from beartype import beartype
+
+from jaxtyping import Float, Bool, Int, Int64, jaxtyped
 
 import torch
 
 
 class ExperienceBuffer:
-    def __init__(self, num_workers: int, device: torch.device) -> None:
+    def __init__(self, num_workers: int, device: torch.device, size: int | None = None) -> None:
         self.num_workers = num_workers
 
-        self.states: List[torch.Tensor] = []
-        self.actions: List[torch.Tensor] = []
-        self.prev_actions: List[torch.Tensor] = [
-            torch.zeros(size=(self.num_workers,), dtype=torch.int64, device=device)
-        ]
-        self.values: List[torch.Tensor] = []
-        self.rewards: List[torch.Tensor] = []
-        self.dones: List[torch.Tensor] = []
-        self.log_probs: List[torch.Tensor] = []
+        self.states: deque[Float[torch.Tensor, "worker channels height width"]] = deque(maxlen=size)
+        self.actions: deque[Int64[torch.Tensor, "worker"]] = deque(maxlen=size)
+        self.prev_actions: deque[Int64[torch.Tensor, "worker"]] = deque(maxlen=size + 1)
+        self.prev_actions.append(torch.zeros(size=(self.num_workers,), dtype=torch.int64, device=device))
+        self.values: deque[Float[torch.Tensor, "worker"]] = deque(maxlen=size)
+        self.rewards: deque[Float[torch.Tensor, "worker"]] = deque(maxlen=size)
+        self.dones: deque[Int64[torch.Tensor, "worker"]] = deque(maxlen=size)
+        self.log_probs: deque[Float[torch.Tensor, "worker"]] = deque(maxlen=size)
 
         self.device = device
 
-    def buffer_states(self, states: torch.Tensor) -> None:
-        assert states.dim() == 4
-        assert states.size(0) == self.num_workers
+    def buffer_states(self, states: Float[torch.Tensor, "worker channels height width"]) -> None:
         self.states.append(states.to(torch.float32))
 
-    def buffer_actions(self, actions: torch.Tensor) -> None:
-        assert actions.dim() == 1
-        assert actions.size(0) == self.num_workers
+    def buffer_actions(self, actions: Int[torch.Tensor, "worker"]) -> None:
         actions_cast = actions.to(torch.int64)
         self.prev_actions.append(actions_cast)
         self.actions.append(actions_cast)
 
-    def buffer_values(self, values: torch.Tensor) -> None:
-        assert values.dim() == 1
-        assert values.size(0) == self.num_workers
+    def buffer_values(self, values: Float[torch.Tensor, "worker"]) -> None:
         self.values.append(values.to(torch.float32))
 
-    def buffer_log_probs(self, log_probs: torch.Tensor) -> None:
-        assert log_probs.dim() == 1
-        assert log_probs.size(0) == self.num_workers
+    def buffer_log_probs(self, log_probs: Float[torch.Tensor, "worker"]) -> None:
         self.log_probs.append(log_probs.to(torch.float32))
 
-    def buffer_rewards(self, rewards: torch.Tensor) -> None:
-        assert rewards.dim() == 1
-        assert rewards.size(0) == self.num_workers
+    def buffer_rewards(self, rewards: Float[torch.Tensor, "worker"]) -> None:
         self.rewards.append(rewards.to(torch.float32))
 
-    def buffer_dones(self, dones: torch.Tensor) -> None:
-        assert dones.dim() == 1
-        assert dones.size(0) == self.num_workers
+    def buffer_dones(self, dones: Bool[torch.Tensor, "worker"]) -> None:
         dones_cast = dones.to(torch.int64)
-        self.prev_actions.append(torch.multiply(self.prev_actions.pop(-1), 1 - dones_cast))
+        self.prev_actions.append(torch.multiply(self.prev_actions.pop(), 1 - dones_cast))
         self.dones.append(dones_cast)
 
     def reset(self, forget_prev_action: bool = False) -> None:
-        self.states = []
+        self.states.clear()
         if forget_prev_action:
-            self.prev_actions = [torch.zeros(size=(self.num_workers,), dtype=torch.int64, device=self.device)]
+            self.prev_actions.clear()
+            self.prev_actions.append(torch.zeros(size=(self.num_workers,), dtype=torch.int64, device=self.device))
         else:
-            self.prev_actions = self.prev_actions[-1:]
-        self.dones = []
-        self.actions = []
-        self.values = []
-        self.rewards = []
-        self.log_probs = []
-
-    def get_last_states(self) -> torch.Tensor:
+            final_action = self.prev_actions[-1]
+            self.prev_actions.clear()
+            self.prev_actions.append(final_action)
+        self.dones.clear()
+        self.actions.clear()
+        self.values.clear()
+        self.rewards.clear()
+        self.log_probs.clear()
+
+    def get_last_states(self) -> Float[torch.Tensor, "worker channels height width"]:
         return self.states[-1]
 
-    def get_last_actions(self) -> torch.Tensor:
+    def get_last_actions(self) -> Float[torch.Tensor, "worker"]:
         return self.prev_actions[-1]
 
-    def get_state_buffer(self) -> torch.Tensor:
-        return torch.stack(self.states)
+    def get_state_buffer(self) -> Float[torch.Tensor, "buffer worker channels height width"]:
+        return torch.stack(list(self.states))
 
-    def get_action_buffer(self) -> torch.Tensor:
-        return torch.stack(self.actions)
+    def get_action_buffer(self) -> Int64[torch.Tensor, "buffer worker"]:
+        return torch.stack(list(self.actions))
 
-    def get_prev_action_buffer(self) -> torch.Tensor:
-        return torch.stack(self.prev_actions[:-1])
+    def get_prev_action_buffer(self) -> Int64[torch.Tensor, "buffer worker"]:
+        return torch.stack(list(self.prev_actions)[:-1])
 
-    def get_value_buffer(self) -> torch.Tensor:
-        return torch.stack(self.values)
+    def get_value_buffer(self) -> Float[torch.Tensor, "buffer worker"]:
+        return torch.stack(list(self.values))
 
-    def get_log_prob_buffer(self) -> torch.Tensor:
-        return torch.stack(self.log_probs)
+    def get_log_prob_buffer(self) -> Float[torch.Tensor, "buffer worker"]:
+        return torch.stack(list(self.log_probs))
 
-    def get_reward_buffer(self) -> torch.Tensor:
-        return torch.stack(self.rewards)
+    def get_reward_buffer(self) -> Float[torch.Tensor, "buffer worker"]:
+        return torch.stack(list(self.rewards))
 
-    def get_dones_buffer(self) -> torch.Tensor:
-        return torch.stack(self.dones)
+    def get_dones_buffer(self) -> Int64[torch.Tensor, "buffer worker"]:
+        return torch.stack(list(self.dones))

src/agents/ppo_agent.py

@@ -51,10 +51,14 @@ def __init__(
             device_name = torch.cuda.get_device_name(self.cpu)
             log.info(f"Using CPU {device_name}")
 
-        self.experience_buffer = ExperienceBuffer(self.num_workers, device=self.device)
+        self.experience_buffer = ExperienceBuffer(
+            self.num_workers, size=self.config.exp_buffer_size, device=self.device
+        )
 
         self.actor_critic = get_model(config=self.config.model, env_info=self.env_info).to(self.device)
 
+        self.sampling_strategy = self.config.sampling_strategy
+
         mlflow.log_text(
             str(
                 torchinfo.summary(
@@ -106,10 +110,10 @@ def _compute_probs_values(self) -> Tuple[torch.Tensor, torch.Tensor]:
     def next_actions(self, train: bool = True) -> Tuple[List[int], List[float]]:
         with torch.no_grad():
             probs, values = self._compute_probs_values()
+            actions = self.sampling_strategy.sample_action(probs)
 
-            action_dist = torch.distributions.Categorical(probs=probs)
-            actions = action_dist.sample()
-            log_probs = action_dist.log_prob(actions)
+            eps = torch.finfo(probs.dtype).eps
+            log_probs = torch.log(probs.clamp(min=eps, max=1 - eps)).gather(-1, actions.unsqueeze(-1)).squeeze(-1)
         self.experience_buffer.buffer_values(values)
         self.experience_buffer.buffer_log_probs(log_probs)
         self.experience_buffer.buffer_actions(actions)
@@ -236,10 +240,10 @@ def _current_entropy_loss_weight(self) -> float:
 
     def update(self) -> None:
         losses = {
-            "actor": 0.0,
-            "critic": 0.0,
-            "entropy": 0.0,
-            "total": 0.0,
+            "actor_loss": 0.0,
+            "critic_loss": 0.0,
+            "entropy_loss": 0.0,
+            "total_loss": 0.0,
         }
 
         dataset = self._create_dataset_from_buffers()
@@ -261,19 +265,25 @@ def update(self) -> None:
 
                 # Actor loss
                 act_loss, action_dist = self._calculate_actor_loss(probs, b_actions, b_log_probs, b_advantages)
-                losses["actor"] += act_loss.item()
+                losses["actor_loss"] += act_loss.item()
 
                 # Critic loss
                 crit_loss = self._calculate_critic_loss(v, b_returns, b_values)
-                losses["critic"] += crit_loss.item()
+                losses["critic_loss"] += crit_loss.item()
 
                 # Entropy loss
                 entropy = torch.mean(action_dist.entropy())
-                losses["entropy"] += entropy.item()
+                losses["entropy_loss"] += entropy.item()
 
                 # Total
-                loss = act_loss + self.critic_loss_weight * crit_loss - self._current_entropy_loss_weight() * entropy
-                losses["total"] += loss.item()
+                loss = torch.add(
+                    act_loss,
+                    torch.sub(
+                        torch.mul(self.critic_loss_weight, crit_loss),
+                        torch.mul(self._current_entropy_loss_weight(), entropy),
+                    ),
+                )
+                losses["total_loss"] += loss.item()
                 total_epoch_loss += loss.item()
                 total_epoch_batches += 1
 
@@ -287,6 +297,7 @@ def update(self) -> None:
             losses[key] /= len(dataloader) * self.epochs_per_update
         log.debug(f"Update finished. Losses: {losses}")
         self.scheduler.step()
+        self.sampling_strategy.update(-losses["total_loss"])
 
         mlflow.log_metrics(losses, self.update_step)
         self.update_step += 1

src/config/agent.py

@@ -2,6 +2,7 @@
 
 from .learning_rate import LRConfig
 from .model import ModelConfig
+from .sampling_strategy import SamplingStrategy
 
 
 @dataclass
@@ -18,8 +19,10 @@ class RandomAgentConfig(AgentConfig):
 class PPOAgentConfig(AgentConfig):
     model: ModelConfig
     learning_rate: LRConfig
+    sampling_strategy: SamplingStrategy
     gamma: float
     tau: float
+    exp_buffer_size: int
     epochs_per_update: int
     total_updates: int
     batch_size: int
@@ -32,6 +35,7 @@ class PPOAgentConfig(AgentConfig):
     def __post_init__(self) -> None:
         assert 0 <= self.gamma <= 1
         assert 0 <= self.tau <= 1
+        assert self.exp_buffer_size >= 1
         assert self.epochs_per_update >= 1
         assert self.total_updates >= 1
         assert self.batch_size >= 1