homa 0.2.95__py3-none-any.whl → 0.3.2__py3-none-any.whl

homa/rl/buffers/concerns/HasRecordAlternatives.py

@@ -0,0 +1,12 @@
+class HasRecordAlternatives:
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def add(self, *args, **kwargs) -> None:
+        self.record(*args, **kwargs)
+
+    def push(self, *args, **kwargs) -> None:
+        self.record(*args, **kwargs)
+
+    def append(self, *args, **kwargs) -> None:
+        self.record(*args, **kwargs)

homa/rl/buffers/concerns/ResetsCollection.py

@@ -0,0 +1,9 @@
+from collections import deque
+
+
+class ResetsCollection:
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def reset(self):
+        self.collection = deque(maxlen=self.capacity)

homa/rl/buffers/concerns/__init__.py

@@ -0,0 +1,2 @@
+from .HasRecordAlternatives import HasRecordAlternatives
+from .ResetsCollection import ResetsCollection

homa/rl/diayn/Actor.py

@@ -0,0 +1,54 @@
+import torch
+from torch.distributions import Normal
+from .modules import ContinuousActorModule
+from ...core.concerns import MovesNetworkToDevice
+
+
+class Actor(MovesNetworkToDevice):
+    def __init__(
+        self,
+        state_dimension: int,
+        action_dimension: int,
+        num_skills: int,
+        hidden_dimension: int,
+        lr: float,
+        decay: float,
+        epsilon: float,
+        min_std: float,
+        max_std: float,
+    ):
+        self.epsilon: float = epsilon
+        self.network = ContinuousActorModule(
+            state_dimension=state_dimension,
+            action_dimension=action_dimension,
+            hidden_dimension=hidden_dimension,
+            num_skills=num_skills,
+            min_std=min_std,
+            max_std=max_std,
+        )
+        self.optimizer = torch.optim.AdamW(
+            self.network.parameters(), lr=lr, weight_decay=decay
+        )
+
+    def action(self, state: torch.Tensor, skill: torch.Tensor):
+        mean, std = self.network(state, skill)
+        std = std.exp()
+        distribution = Normal(mean, std)
+        raw_action = distribution.rsample()
+        action = torch.tanh(raw_action)
+        corrected_probabilities = torch.log(1.0 - action.pow(2) + self.epsilon)
+        probabilities = distribution.log_prob(raw_action) - corrected_probabilities
+        probabilities = probabilities.sum(dim=-1, keepdim=True)
+        return action, probabilities
+
+    def train(self, advantages: torch.Tensor, probabilities: torch.Tensor) -> float:
+        self.optimizer.zero_grad()
+        loss = self.loss(advantages=advantages, probabilities=probabilities)
+        loss.backward()
+        self.optimizer.step()
+        return loss.item()
+
+    def loss(
+        self, advantages: torch.Tensor, probabilities: torch.Tensor
+    ) -> torch.Tensor:
+        return -(probabilities * advantages.detach()).mean()

homa/rl/diayn/Critic.py

@@ -0,0 +1,41 @@
+import torch
+from .modules import CriticModule
+from ...core.concerns import MovesNetworkToDevice
+
+
+class Critic(MovesNetworkToDevice):
+    def __init__(
+        self,
+        state_dimension: int,
+        hidden_dimension: int,
+        num_skills: int,
+        lr: float,
+        decay: float,
+        gamma: float,
+    ):
+        self.network = CriticModule(
+            state_dimension=state_dimension,
+            hidden_dimension=hidden_dimension,
+            num_skills=num_skills,
+        )
+        self.optimizer = torch.optim.AdamW(
+            self.network.parameters(), lr=lr, weight_decay=decay
+        )
+        self.criterion = torch.nn.SmoothL1Loss()
+        self.gamma: float = gamma
+
+    def train(self, advantages: torch.Tensor):
+        self.optimizer.zero_grad()
+        loss = self.loss(advantages=advantages)
+        loss.backward()
+        self.optimizer.step()
+
+    def loss(self, advantages: torch.Tensor):
+        return advantages.pow(2).mean()
+
+    def values(self, states: torch.Tensor, skills: torch.Tensor):
+        return self.network(states, skills)
+
+    @torch.no_grad()
+    def values_(self, *args, **kwargs):
+        return self.values(*args, **kwargs)

homa/rl/diayn/Discriminator.py

@@ -0,0 +1,45 @@
+import torch
+import numpy
+from .modules import DiscriminatorModule
+from ...core.concerns import MovesNetworkToDevice
+
+
+class Discriminator(MovesNetworkToDevice):
+    def __init__(
+        self,
+        state_dimension: int,
+        hidden_dimension: int,
+        num_skills: int,
+        decay: float,
+        lr: float,
+    ):
+        self.num_skills: int = num_skills
+        self.network = DiscriminatorModule(
+            state_dimension=state_dimension,
+            hidden_dimension=hidden_dimension,
+            num_skills=num_skills,
+        )
+        self.optimizer = torch.optim.AdamW(
+            self.network.parameters(), lr=lr, weight_decay=decay
+        )
+        self.criterion = torch.nn.CrossEntropyLoss()
+
+    def loss(self, states: torch.Tensor, skills_indices: torch.Tensor):
+        logits = self.network(states)
+        return self.criterion(logits, skills_indices)
+
+    @torch.no_grad()
+    def reward(self, state: torch.Tensor, skill_index: torch.Tensor):
+        logits = self.network(state)
+        probabilities = torch.nn.functional.log_softmax(logits, dim=-1)
+        entropy = numpy.log(1.0 / self.num_skills)
+        if skill_index.dim() == 1:
+            skill_index = skill_index.unsqueeze(-1)
+        reward = probabilities.gather(1, skill_index.long()) - entropy
+        return reward.squeeze(-1)
+
+    def train(self, states: torch.Tensor, skills_indices: torch.Tensor):
+        self.optimizer.zero_grad()
+        loss = self.loss(states=states, skills_indices=skills_indices)
+        loss.backward()
+        self.optimizer.step()

homa/rl/diayn/__init__.py

@@ -0,0 +1,3 @@
+from .Actor import Actor
+from .Critic import Critic
+from .Discriminator import Discriminator

homa/rl/diayn/modules/ContinuousActorModule.py

@@ -0,0 +1,42 @@
+import torch
+
+
+class ContinuousActorModule(torch.nn.Module):
+    def __init__(
+        self,
+        state_dimension: int,
+        action_dimension: int,
+        hidden_dimension: int,
+        num_skills: int,
+        min_std: float,
+        max_std: float,
+    ):
+        super().__init__()
+        self.state_dimension: int = state_dimension
+        self.action_dimension: int = action_dimension
+        self.num_skills: int = num_skills
+        self.hidden_dimension: int = hidden_dimension
+        self.input_dimension: int = self.state_dimension + self.num_skills
+        self.min_std: float = min_std
+        self.max_std: float = max_std
+
+        self.phi = torch.nn.Sequential(
+            torch.nn.Linear(self.input_dimension, self.hidden_dimension),
+            torch.nn.ReLU(),
+            torch.nn.Linear(self.hidden_dimension, self.hidden_dimension),
+            torch.nn.ReLU(),
+            torch.nn.Linear(self.hidden_dimension, self.hidden_dimension),
+        )
+
+        self.mu = torch.nn.Linear(self.hidden_dimension, self.action_dimension)
+        self.xi = torch.nn.Linear(self.hidden_dimension, self.action_dimension)
+
+    def forward(self, state: torch.Tensor, skill: torch.Tensor) -> torch.Tensor:
+        # fix the size to be one state per batch
+        state = state.view(state.size(0), -1)
+
+        psi = torch.cat([state, skill], dim=-1)
+        features = self.phi(psi)
+        mean = self.mu(features)
+        std = self.xi(features).clamp(self.min_std, self.max_std)
+        return mean, std

homa/rl/diayn/modules/CriticModule.py

@@ -0,0 +1,28 @@
+import torch
+
+
+class CriticModule(torch.nn.Module):
+    def __init__(
+        self,
+        state_dimension: int,
+        hidden_dimension: int,
+        num_skills: int,
+    ):
+        super().__init__()
+        self.state_dimension: int = state_dimension
+        self.num_skills: int = num_skills
+        self.hidden_dimension: int = hidden_dimension
+        self.input_dimension: int = self.state_dimension + self.num_skills
+
+        self.phi = torch.nn.Sequential(
+            torch.nn.Linear(self.input_dimension, self.hidden_dimension),
+            torch.nn.ReLU(),
+            torch.nn.Linear(self.hidden_dimension, self.hidden_dimension),
+            torch.nn.ReLU(),
+        )
+        self.fc = (torch.nn.Linear(self.hidden_dimension, 1),)
+
+    def forward(self, state: torch.Tensor, skill: torch.Tensor) -> torch.Tensor:
+        psi = torch.cat([state, skill], dim=-1)
+        features = self.phi(psi)
+        return self.fc(features).squeeze(-1)

homa/rl/diayn/modules/DiscriminatorModule.py

@@ -0,0 +1,24 @@
+import torch
+from typing import Type
+
+
+class DiscriminatorModule(torch.nn.Module):
+    def __init__(self, state_dimension: int, hidden_dimension: int, num_skills: int):
+        super().__init__()
+        self.state_dimension: int = state_dimension
+        self.hidden_dimension: int = hidden_dimension
+        self.num_skills: int = num_skills
+
+        self.phi: Type[torch.nn.Sequential] = torch.nn.Sequential(
+            torch.nn.Linear(self.state_dimension, self.hidden_dimension),
+            torch.nn.ReLU(),
+            torch.nn.Linear(self.hidden_dimension, self.hidden_dimension),
+            torch.nn.ReLU(),
+        )
+        self.fc: Type[torch.nn.Linear] = torch.nn.Linear(
+            self.hidden_dimension, self.num_skills
+        )
+
+    def forward(self, state: torch.Tensor) -> torch.Tensor:
+        features: torch.Tensor = self.phi(state)
+        return self.fc(features)

homa/rl/diayn/modules/__init__.py

@@ -0,0 +1,3 @@
+from .DiscriminatorModule import DiscriminatorModule
+from .CriticModule import CriticModule
+from .ContinuousActorModule import ContinuousActorModule

homa/rl/sac/SoftActor.py

@@ -0,0 +1,70 @@
+import torch
+import numpy
+from .SoftCritic import SoftCritic
+from .modules import SoftActorModule
+from ...core.concerns import MovesNetworkToDevice
+
+
+class SoftActor(MovesNetworkToDevice):
+    def __init__(
+        self,
+        state_dimension: int,
+        hidden_dimension: int,
+        action_dimension: int,
+        lr: float,
+        weight_decay: float,
+        alpha: float,
+        min_std: float,
+        max_std: float,
+    ):
+        self.alpha: float = alpha
+
+        self.network = SoftActorModule(
+            state_dimension=state_dimension,
+            hidden_dimension=hidden_dimension,
+            action_dimension=action_dimension,
+            min_std=min_std,
+            max_std=max_std,
+        )
+        self.optimizer = torch.optim.AdamW(
+            self.network.parameters(), lr=lr, weight_decay=weight_decay
+        )
+
+    def train(self, states: torch.Tensor, critic: SoftCritic):
+        self.network.train()
+        self.optimizer.zero_grad()
+        loss = self.loss(states=states, critic=critic)
+        loss.backward()
+        self.optimizer.step()
+
+    def loss(self, states: torch.Tensor, critic: SoftCritic) -> torch.Tensor:
+        actions, probabilities = self.sample(states)
+        q_alpha, q_beta = critic.network(states, actions)
+        q = torch.min(q_alpha, q_beta)
+        return (self.alpha * probabilities - q).mean()
+
+    def process_state(self, state: numpy.ndarray | torch.Tensor) -> torch.Tensor:
+        if isinstance(state, numpy.ndarray):
+            state = torch.from_numpy(state).float()
+
+        if state.ndim < 2:
+            state = state.unsqueeze(0)
+
+        return state
+
+    def sample(self, state: numpy.ndarray | torch.Tensor):
+        state = self.process_state(state)
+
+        mean, std = self.network(state)
+        # following line prevents standard deviations to be negative
+        std = std.exp()
+
+        distribution = torch.distributions.Normal(mean, std)
+
+        pre_tanh = distribution.rsample()
+        action = torch.tanh(pre_tanh)
+
+        probabilities = distribution.log_prob(pre_tanh).sum(dim=1, keepdim=True)
+        probabilities -= torch.log(1 - action.pow(2) + 1e-6).sum(dim=1, keepdim=True)
+
+        return action, probabilities

homa/rl/sac/SoftCritic.py

@@ -0,0 +1,98 @@
+import torch
+from torch.nn.functional import mse_loss as mse
+from .modules import DualSoftCriticModule
+from .SoftActor import SoftActor
+from ..utils import soft_update
+from ...core.concerns import MovesNetworkToDevice
+
+
+class SoftCritic(MovesNetworkToDevice):
+    def __init__(
+        self,
+        state_dimension: int,
+        hidden_dimension: int,
+        action_dimension: int,
+        lr: float,
+        weight_decay: float,
+        gamma: float,
+        alpha: float,
+    ):
+        self.gamma: float = gamma
+        self.alpha: float = alpha
+
+        self.network = DualSoftCriticModule(
+            state_dimension=state_dimension,
+            hidden_dimension=hidden_dimension,
+            action_dimension=action_dimension,
+        )
+        self.target = DualSoftCriticModule(
+            state_dimension=state_dimension,
+            hidden_dimension=hidden_dimension,
+            action_dimension=action_dimension,
+        )
+
+        # copy source to target when initiated
+        self.target.load_state_dict(self.network.state_dict())
+
+        self.optimizer = torch.optim.AdamW(
+            self.network.parameters(), lr=lr, weight_decay=weight_decay
+        )
+
+    def train(
+        self,
+        states: torch.Tensor,
+        actions: torch.Tensor,
+        rewards: torch.Tensor,
+        terminations: torch.Tensor,
+        next_states: torch.Tensor,
+        actor: SoftActor,
+    ):
+        self.network.train()
+        self.optimizer.zero_grad()
+        loss = self.loss(
+            states=states,
+            actions=actions,
+            rewards=rewards,
+            terminations=terminations,
+            next_states=next_states,
+            actor=actor,
+        )
+        loss.backward()
+        self.optimizer.step()
+
+    def loss(
+        self,
+        states: torch.Tensor,
+        actions: torch.Tensor,
+        rewards: torch.Tensor,
+        terminations: torch.Tensor,
+        next_states: torch.Tensor,
+        actor: torch.nn.Module,
+    ):
+        q_alpha, q_beta = self.network(states, actions)
+        target = self.calculate_target(
+            rewards=rewards,
+            terminations=terminations,
+            next_states=next_states,
+            actor=actor,
+        )
+        return mse(q_alpha, target) + mse(q_beta, target)
+
+    @torch.no_grad()
+    def calculate_target(
+        self,
+        rewards: torch.Tensor,
+        terminations: torch.Tensor,
+        next_states: torch.Tensor,
+        actor: SoftActor,
+    ):
+        termination_mask = 1 - terminations
+        next_actions, next_probabilities = actor.sample(next_states)
+        q_alpha, q_beta = self.target(next_states, next_actions)
+        q = torch.min(q_alpha, q_beta)
+        entropy_q = q - self.alpha * next_probabilities
+        return rewards + self.gamma * termination_mask * entropy_q
+
+    def update(self, tau: float):
+        soft_update(network=self.network.alpha, target=self.target.alpha, tau=tau)
+        soft_update(network=self.network.beta, target=self.target.beta, tau=tau)

homa/rl/sac/__init__.py

@@ -0,0 +1,2 @@
+from .SoftActor import SoftActor
+from .SoftCritic import SoftCritic

homa/rl/sac/modules/DualSoftCriticModule.py

@@ -0,0 +1,22 @@
+import torch
+from .SoftCriticModule import SoftCriticModule
+
+
+class DualSoftCriticModule(torch.nn.Module):
+    def __init__(
+        self, state_dimension: int, hidden_dimension: int, action_dimension: int
+    ):
+        super().__init__()
+        self.alpha = SoftCriticModule(
+            state_dimension=state_dimension,
+            hidden_dimension=hidden_dimension,
+            action_dimension=action_dimension,
+        )
+        self.beta = SoftCriticModule(
+            state_dimension=state_dimension,
+            hidden_dimension=hidden_dimension,
+            action_dimension=action_dimension,
+        )
+
+    def forward(self, state: torch.Tensor, action: torch.Tensor):
+        return self.alpha(state, action), self.beta(state, action)

homa/rl/sac/modules/SoftActorModule.py

@@ -0,0 +1,35 @@
+import torch
+
+
+class SoftActorModule(torch.nn.Module):
+    def __init__(
+        self,
+        state_dimension: int,
+        hidden_dimension: int,
+        action_dimension: int,
+        min_std: float,
+        max_std: float,
+    ):
+        super().__init__()
+
+        self.state_dimension: int = state_dimension
+        self.hidden_dimension: int = hidden_dimension
+        self.action_dimension: int = action_dimension
+        self.min_std: float = float(min_std)
+        self.max_std: float = float(max_std)
+
+        self.phi = torch.nn.Sequential(
+            torch.nn.Linear(self.state_dimension, self.hidden_dimension),
+            torch.nn.ReLU(),
+            torch.nn.Linear(self.hidden_dimension, self.hidden_dimension),
+            torch.nn.ReLU(),
+        )
+        self.mu = torch.nn.Linear(self.hidden_dimension, self.action_dimension)
+        self.xi = torch.nn.Linear(self.hidden_dimension, self.action_dimension)
+
+    def forward(self, state: torch.Tensor):
+        features = self.phi(state)
+        mean = self.mu(features)
+        std = self.xi(features)
+        std = std.clamp(self.min_std, self.max_std)
+        return mean, std

homa/rl/sac/modules/SoftCriticModule.py

@@ -0,0 +1,30 @@
+import torch
+
+
+class SoftCriticModule(torch.nn.Module):
+    def __init__(
+        self,
+        state_dimension: int,
+        hidden_dimension: int,
+        action_dimension: int,
+    ):
+        super().__init__()
+
+        self.state_dimension: int = state_dimension
+        self.action_dimension: int = action_dimension
+        self.hidden_dimension: int = hidden_dimension
+
+        self.phi = torch.nn.Sequential(
+            torch.nn.Linear(
+                self.state_dimension + self.action_dimension, self.hidden_dimension
+            ),
+            torch.nn.ReLU(),
+            torch.nn.Linear(self.hidden_dimension, self.hidden_dimension),
+            torch.nn.ReLU(),
+        )
+        self.fc = torch.nn.Linear(self.hidden_dimension, 1)
+
+    def forward(self, state: torch.Tensor, action: torch.Tensor):
+        psi = torch.cat([state, action], dim=1)
+        features = self.phi(psi)
+        return self.fc(features)

homa/rl/sac/modules/__init__.py

@@ -0,0 +1,3 @@
+from .SoftActorModule import SoftActorModule
+from .SoftCriticModule import SoftCriticModule
+from .DualSoftCriticModule import DualSoftCriticModule

homa/rl/utils.py

@@ -0,0 +1,7 @@
+import torch
+
+
+@torch.no_grad()
+def soft_update(network: torch.nn.Module, target: torch.nn.Module, tau: float):
+    for s, t in zip(network.parameters(), target.parameters()):
+        t.data.copy_(tau * s.data + (1 - tau) * t.data)

homa/vision/Resnet.py

@@ -2,12 +2,12 @@ import torch
 from .modules import ResnetModule
 from .Classifier import Classifier
 from .concerns import Trainable, ReportsMetrics
-from ..device import get_device
+from ..core.concerns import MovesNetworkToDevice
 
 
-class Resnet(Classifier, Trainable, ReportsMetrics):
+class Resnet(Classifier, Trainable, ReportsMetrics, MovesNetworkToDevice):
     def __init__(self, num_classes: int, lr: float = 0.001):
         super().__init__()
-        self.network = ResnetModule(num_classes).to(get_device())
+        self.network = ResnetModule(num_classes)
         self.criterion = torch.nn.CrossEntropyLoss()
         self.optimizer = torch.optim.SGD(self.network.parameters(), lr=lr, momentum=0.9)

homa/vision/Swin.py

@@ -2,12 +2,24 @@ import torch
 from .Classifier import Classifier
 from .concerns import Trainable, ReportsMetrics
 from .modules import SwinModule
-from ..device import get_device
+from ..core.concerns import MovesNetworkToDevice
 
 
-class Swin(Classifier, Trainable, ReportsMetrics):
-    def __init__(self, num_classes: int, lr: float = 0.0001):
+class Swin(Classifier, Trainable, ReportsMetrics, MovesNetworkToDevice):
+    def __init__(
+        self,
+        num_classes: int,
+        lr: float = 0.0001,
+        decay: float = 0.0,
+        variant: str = "base",
+        weights="DEFAULT",
+    ):
         super().__init__()
-        self.network = SwinModule(num_classes=num_classes).to(get_device())
-        self.optimizer = torch.optim.AdamW(self.network.parameters(), lr=lr)
+        self.num_classes = num_classes
+        self.network = SwinModule(
+            num_classes=self.num_classes, variant=variant, weights=weights
+        )
+        self.optimizer = torch.optim.AdamW(
+            self.network.parameters(), lr=lr, weight_decay=decay
+        )
         self.criterion = torch.nn.CrossEntropyLoss()

homa/vision/modules/SwinModule.py

@@ -1,21 +1,29 @@
 import torch
-from torchvision.models import swin_v2_b
+from torchvision.models import swin_v2_b, swin_v2_s, swin_v2_t
 from torch.nn.init import kaiming_uniform_ as kaiming
 
 
 class SwinModule(torch.nn.Module):
-    def __init__(self, num_classes: int):
+    def __init__(self, num_classes: int, variant: str, weights):
         super().__init__()
-        self.num_classes = num_classes
-        self._create_encoder()
-        self._create_fc()
+        self._create_encoder(variant=variant, weights=weights)
+        self._create_fc(num_classes=num_classes)
 
-    def _create_encoder(self):
-        self.encoder = swin_v2_b(weights="DEFAULT")
+    def variant_instance(self, variant: str):
+        variant_map = {"tiny": swin_v2_t, "small": swin_v2_s, "base": swin_v2_b}
+        return variant_map.get(variant)
+
+    def _create_encoder(self, variant: str, weights):
+        if variant not in ["tiny", "small", "base"]:
+            raise ValueError(
+                f"Swin variant needs to be one of [tiny, small, base]. Invalid {variant} was provided."
+            )
+        instance = self.variant_instnace(variant)
+        self.encoder = instance(weights=weights)
         self.encoder.head = torch.nn.Identity()
 
-    def _create_fc(self):
-        self.fc = torch.nn.Linear(1024, self.num_classes)
+    def _create_fc(self, num_classes: int):
+        self.fc = torch.nn.Linear(1024, num_classes)
         kaiming(self.fc.weight, mode="fan_in", nonlinearity="relu")
 
     def forward(self, images: torch.Tensor):

{homa-0.2.95.dist-info → homa-0.3.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: homa
-Version: 0.2.95
+Version: 0.3.2
 Summary: A curated list of machine learning and deep learning helpers.
 Author-email: Taha Shieenavaz <tahashieenavaz@gmail.com>
 Requires-Python: >=3.7