Merge pull request #11 from LondonNode/feature/optimizations

Feature/optimizations

Author: 0xprofessooor

pearll/buffers/base_buffer.py

@@ -4,12 +4,14 @@
 
 import numpy as np
 import psutil
+import torch as T
 from gym import Env
 from gym.vector import VectorEnv
 
+from pearll import settings
 from pearll.common.enumerations import TrajectoryType
 from pearll.common.type_aliases import Observation, Trajectories
-from pearll.common.utils import get_space_shape, to_torch
+from pearll.common.utils import get_space_shape
 
 
 class BaseBuffer(ABC):
@@ -128,13 +130,15 @@ def _transform_samples(
 
         # return torch tensors instead of numpy arrays
         if dtype == TrajectoryType.TORCH:
-            observations, actions, rewards, next_observations, dones = to_torch(
-                observations,
-                actions,
-                rewards,
-                next_observations,
-                dones,
+            observations = T.from_numpy(observations).to(
+                settings.DEVICE, non_blocking=True
             )
+            actions = T.from_numpy(actions).to(settings.DEVICE, non_blocking=True)
+            rewards = T.from_numpy(rewards).to(settings.DEVICE, non_blocking=True)
+            next_observations = T.from_numpy(next_observations).to(
+                settings.DEVICE, non_blocking=True
+            )
+            dones = T.from_numpy(dones).to(settings.DEVICE, non_blocking=True)
 
         return Trajectories(
             observations=observations,

pearll/buffers/her_buffer.py

@@ -1,8 +1,10 @@
 from typing import Dict, Tuple, Union
 
 import numpy as np
+import torch as T
 from gym.core import GoalEnv
 
+from pearll import settings
 from pearll.buffers.base_buffer import BaseBuffer
 from pearll.common.enumerations import GoalSelectionStrategy, TrajectoryType
 from pearll.common.type_aliases import DictTrajectories, Tensor
@@ -219,7 +221,31 @@ def sample(
         else:
             batch_inds = np.random.randint(0, end_idx, size=batch_size)
 
-        trajectories = self._sample_trajectories(batch_inds)
+        trajectories = list(self._sample_trajectories(batch_inds))
+
+        if dtype == TrajectoryType.TORCH:
+            trajectories[0]["observation"] = T.from_numpy(
+                trajectories[0]["observation"]
+            ).to(settings.DEVICE, non_blocking=True)
+            trajectories[0]["desired_goal"] = T.from_numpy(
+                trajectories[0]["desired_goal"]
+            ).to(settings.DEVICE, non_blocking=True)
+            trajectories[1] = T.from_numpy(trajectories[1]).to(
+                settings.DEVICE, non_blocking=True
+            )
+            trajectories[2] = T.from_numpy(trajectories[2]).to(
+                settings.DEVICE, non_blocking=True
+            )
+            trajectories[3]["observation"] = T.from_numpy(
+                trajectories[3]["observation"]
+            ).to(settings.DEVICE, non_blocking=True)
+            trajectories[3]["desired_goal"] = T.from_numpy(
+                trajectories[3]["desired_goal"]
+            ).to(settings.DEVICE, non_blocking=True)
+            trajectories[4] = T.from_numpy(trajectories[4]).to(
+                settings.DEVICE, non_blocking=True
+            )
+
         return DictTrajectories(
             observations=trajectories[0],
             actions=trajectories[1],

pearll/common/utils.py

@@ -56,7 +56,7 @@ def to_numpy(*data) -> Union[Tuple[np.ndarray], np.ndarray]:
         if isinstance(el, T.Tensor):
             result[i] = el.detach().cpu().numpy()
         else:
-            result[i] = np.array(el)
+            result[i] = np.asarray(el)
 
     if len(data) == 1:
         return result[0]

pearll/models/actor_critics.py

@@ -8,7 +8,7 @@
 from pearll import settings
 from pearll.common.enumerations import Distribution
 from pearll.common.type_aliases import Tensor
-from pearll.common.utils import get_space_range, get_space_shape, to_numpy
+from pearll.common.utils import get_space_range, get_space_shape
 from pearll.models.encoders import IdentityEncoder, MLPEncoder
 from pearll.models.heads import (
     BaseActorHead,
@@ -66,7 +66,10 @@ def __init__(
         self.state_info = {}
         self.make_state_info()
         self.state = np.concatenate(
-            [to_numpy(d.flatten()) for d in self.model.state_dict().values()]
+            [
+                d.flatten().detach().cpu().numpy()
+                for d in self.model.state_dict().values()
+            ]
         )
         self.space = Box(low=-1e6, high=1e6, shape=self.state.shape)
         self.space_shape = get_space_shape(self.space)
@@ -217,7 +220,12 @@ def forward(self, observations: Tensor) -> T.Tensor:
         trigger = T.rand(1).item()
 
         if trigger <= self.epsilon:
-            actions = T.randint(low=0, high=action_size, size=q_values.shape[:-1])
+            actions = T.randint(
+                low=0,
+                high=action_size,
+                size=q_values.shape[:-1],
+                device=settings.DEVICE,
+            )
         else:
             _, actions = T.max(q_values, dim=-1)
 

pearll/models/encoders.py

@@ -6,7 +6,7 @@
 
 from pearll.common.type_aliases import Tensor
 from pearll.common.utils import to_numpy
-from pearll.models.utils import concat_obs_actions
+from pearll.models.utils import preprocess_inputs
 
 
 class IdentityEncoder(T.nn.Module):
@@ -19,8 +19,8 @@ def forward(
         self, observations: Tensor, actions: Optional[Tensor] = None
     ) -> T.Tensor:
         # Some algorithms use both the observations and actions as input (e.g. DDPG for conitnuous Q function)
-        observations = concat_obs_actions(observations, actions)
-        return observations
+        input = preprocess_inputs(observations, actions)
+        return input
 
 
 class FlattenEncoder(T.nn.Module):
@@ -34,8 +34,8 @@ def forward(
     ) -> T.Tensor:
         # Some algorithms use both the observations and actions as input (e.g. DDPG for conitnuous Q function)
         # Make sure observations is a torch tensor, get error if numpy for some reason??
-        observations = concat_obs_actions(observations, actions)
-        return T.flatten(observations)
+        input = preprocess_inputs(observations, actions)
+        return T.flatten(input)
 
 
 class MLPEncoder(T.nn.Module):
@@ -48,8 +48,8 @@ def __init__(self, input_size, output_size):
     def forward(
         self, observations: Tensor, actions: Optional[Tensor] = None
     ) -> T.Tensor:
-        observations = concat_obs_actions(observations, actions)
-        return self.model(observations)
+        input = preprocess_inputs(observations, actions)
+        return self.model(input)
 
 
 class CNNEncoder(T.nn.Module):

pearll/models/utils.py

@@ -2,8 +2,8 @@
 
 import torch as T
 
+from pearll import settings
 from pearll.common.type_aliases import Tensor
-from pearll.common.utils import to_torch
 
 
 def trainable_parameters(model: T.nn.Module) -> list:
@@ -31,12 +31,13 @@ def get_mlp_size(data_shape: Union[int, Tuple[int]]) -> int:
     return data_shape
 
 
-def concat_obs_actions(observations: Tensor, actions: Optional[Tensor]) -> T.Tensor:
+def preprocess_inputs(observations: Tensor, actions: Optional[Tensor]) -> T.Tensor:
+    input = T.as_tensor(observations)
+    if input.dim() == 0:
+        input = input.unsqueeze(0)
     if actions is not None:
-        observations, actions = to_torch(observations, actions)
-        if observations.dim() == 0:
-            observations = observations.unsqueeze(0)
+        actions = T.as_tensor(actions)
         if actions.dim() == 0:
             actions = actions.unsqueeze(0)
-        return T.cat([observations, actions], dim=-1).float()
-    return to_torch(observations).float()
+        input = T.cat([input, actions], dim=-1)
+    return input.float().to(settings.DEVICE, non_blocking=True)

pearll/updaters/evolution.py

@@ -2,6 +2,7 @@
 from typing import Optional
 
 import numpy as np
+import torch as T
 from gym.spaces import Discrete, MultiDiscrete
 from torch.distributions import Normal, kl_divergence
 
@@ -11,7 +12,6 @@
     SelectionFunc,
     UpdaterLog,
 )
-from pearll.common.utils import to_torch
 from pearll.models.actor_critics import ActorCritic
 
 
@@ -86,7 +86,7 @@ def __call__(
         """
         # Snapshot current population dist for kl divergence
         # use copy() to avoid modifying the original
-        old_dist = Normal(to_torch(self.mean.copy()), self.std)
+        old_dist = Normal(T.from_numpy(self.mean.copy()), self.std)
 
         # Main update
         self.mean += learning_rate * optimization_direction
@@ -104,7 +104,7 @@ def __call__(
         self.update_networks(population)
 
         # Calculate Log metrics
-        new_dist = Normal(to_torch(self.mean), self.std)
+        new_dist = Normal(T.from_numpy(self.mean), self.std)
         population_entropy = new_dist.entropy().mean()
         population_kl = kl_divergence(old_dist, new_dist).mean()
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pearll"
-version = "0.4.0"
+version = "0.4.1"
 description = "Adaptable tools to make reinforcement learning and evolutionary computation algorithms"
 license = "MIT"
 authors = ["Rohan Tangri <rohan.tangri@gmail.com>"]

tests/test_her.py

@@ -291,9 +291,11 @@ def test_her_sample(goal_selection_strategy, buffer_size):
     observations[pos] = next_observations[pos - 1]
 
     trajectories = buffer.sample(4)
-    sampled_observations = trajectories.observations["observation"]
-    sampled_next_observations = trajectories.next_observations["observation"]
-    her_sampled_goals = trajectories.observations["desired_goal"]
+    sampled_observations = np.asarray(trajectories.observations["observation"])
+    sampled_next_observations = np.asarray(
+        trajectories.next_observations["observation"]
+    )
+    her_sampled_goals = np.asarray(trajectories.observations["desired_goal"])
     # Check if sampled next observations are actually the next observations
     for i, obs in enumerate(sampled_observations):
         array_idx = np.where(observations == obs)[0]