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

homa/core/concerns/MovesNetworkToDevice.py

@@ -0,0 +1,13 @@
+from ...device import move
+
+
+class MovesNetworkToDevice:
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        if not hasattr(self, "network"):
+            raise RuntimeError(
+                "MovesNetworkToDevice assumes the underlying class has a network property."
+            )
+
+        move(self.network)

homa/core/concerns/__init__.py

@@ -0,0 +1 @@
+from .MovesNetworkToDevice import MovesNetworkToDevice

homa/device.py

@@ -23,3 +23,7 @@ def mps():
 
 def device():
     return get_device()
+
+
+def move(module: torch.nn.Module):
+    module.to(get_device())

homa/ensemble/Ensemble.py

@@ -1,16 +1,18 @@
 from .concerns import (
-    ReportsSize,
+    ReportsEnsembleSize,
     StoresModels,
     ReportsClassificationMetrics,
     PredictsProbabilities,
+    SavesEnsembleModels,
 )
 
 
 class Ensemble(
-    ReportsSize,
+    ReportsEnsembleSize,
     ReportsClassificationMetrics,
     PredictsProbabilities,
     StoresModels,
+    SavesEnsembleModels,
 ):
     def __init__(self):
         super().__init__()

homa/ensemble/concerns/{ReportsSize.py → ReportsEnsembleSize.py}

@@ -1,11 +1,11 @@
-class ReportsSize:
+class ReportsEnsembleSize:
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
     @property
     def size(self):
-        return len(self.models)
+        return len(self.weights)
 
     @property
     def length(self):
-        return len(self.models)
+        return self.size

homa/ensemble/concerns/ReportsLogits.py

@@ -8,7 +8,9 @@ class ReportsLogits:
     def logits(self, x: torch.Tensor) -> torch.Tensor:
         batch_size = x.shape[0]
         logits = torch.zeros((batch_size, self.num_classes))
-        for model in self.models:
+        for factory, weight in zip(self.factories, self.weights):
+            model = factory(num_classes=self.num_classes)
+            model.load_state_dict(weight)
             logits += model(x)
         return logits
 

homa/ensemble/concerns/SavesEnsembleModels.py

@@ -0,0 +1,13 @@
+class SavesEnsembleModels:
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def save(self):
+        self.save_factories()
+        self.save_weights()
+
+    def save_factories(self):
+        pass
+
+    def save_weights(self):
+        pass

homa/ensemble/concerns/StoresModels.py

@@ -1,13 +1,14 @@
 import torch
-import io
-from typing import List
+from typing import List, Type
+from collections import OrderedDict
 from ...vision import Model
 
 
 class StoresModels:
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.models: List[torch.nn.Module] = []
+        self.factories: List[Type[torch.nn.Module]] = []
+        self.weights: List[OrderedDict] = []
 
     def record(self, model: Model | torch.nn.Module):
         model_: torch.nn.Module | None = None
@@ -18,12 +19,8 @@ class StoresModels:
         else:
             raise TypeError("Wrong input to ensemble record")
 
-        device = model_.device
-        buffer = io.BytesIO()
-        torch.save(model_.to("cpu"), buffer)
-        buffer.seek(0)
-        model_ = torch.load(buffer, map_location=device)
-        self.models.append(model_)
+        self.factories.append(model_.__class__)
+        self.weights.append(model_.state_dict())
 
     def push(self, *args, **kwargs):
         self.record(*args, **kwargs)

homa/ensemble/concerns/__init__.py

@@ -5,5 +5,6 @@ from .ReportsEnsembleAccuracy import ReportsEnsembleAccuracy
 from .ReportsEnsembleF1 import ReportsEnsembleF1
 from .ReportsEnsembleKappa import ReportsEnsembleKappa
 from .ReportsLogits import ReportsLogits
-from .ReportsSize import ReportsSize
+from .ReportsEnsembleSize import ReportsEnsembleSize
 from .StoresModels import StoresModels
+from .SavesEnsembleModels import SavesEnsembleModels

homa/ensemble/utils.py

@@ -0,0 +1,9 @@
+import torch
+
+
+def get_model_device(model: torch.nn.Module):
+    try:
+        device = next(model.parameters()).device
+    except StopIteration:
+        device = torch.device("cpu")
+    return device

homa/graph/GraphAttention.py

@@ -0,0 +1,13 @@
+import torch
+from .modules import GraphAttentionModule
+from ..core.concerns import MovesNetworkToDevice
+
+
+class GraphAttention(MovesNetworkToDevice):
+    def __init__(self, lr: float = 0.005, decay: float = 5e-4, dropout: float = 0.6):
+        super().__init__()
+        self.network = GraphAttentionModule()
+        self.optimizer = torch.nn.AdamW(
+            self.network.parameters(), lr=lr, weight_decay=decay
+        )
+        self.criterion = torch.nn.CrossEntropyLoss()

homa/graph/__init__.py

@@ -0,0 +1 @@
+from .GraphAttention import GraphAttention

homa/graph/modules/GraphAttentionHeadModule.py

@@ -0,0 +1,37 @@
+import torch
+
+
+class GraphAttentionHeadModule(torch.nn.Module):
+    def __init__(self, in_features: int, out_features: int, alpha=0.2):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.alpha = alpha
+
+        self.W = torch.nn.Linear(in_features, out_features, bias=False)
+        self.a_1 = torch.nn.Parameter(torch.randn(out_features, 1))
+        self.a_2 = torch.nn.Parameter(torch.randn(out_features, 1))
+
+        self.leaky_relu = torch.nn.LeakyReLU(self.alpha)
+        self.elu = torch.nn.ELU()
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        torch.nn.init.xavier_uniform_(self.W.weight, gain=1.414)
+        torch.nn.init.xavier_uniform_(self.a_1, gain=1.414)
+        torch.nn.init.xavier_uniform_(self.a_2, gain=1.414)
+
+    def forward(self, node_features, adj_matrix):
+        N = node_features.size(0)
+        h_prime = self.W(node_features)
+        s1 = torch.matmul(h_prime, self.a_1)
+        s2 = torch.matmul(h_prime, self.a_2)
+        e = s1 + s2.T
+        e = self.leaky_relu(e)
+        zero_vec = -9e15 * torch.ones_like(e)
+        attention_mask = torch.where(
+            adj_matrix > 0, e, zero_vec.to(node_features.device)
+        )
+        attention_weights = F.softmax(attention_mask, dim=1)
+        h_new = torch.matmul(attention_weights, h_prime)
+        return self.elu(h_new)

homa/graph/modules/MultiHeadGraphAttentionModule.py

@@ -0,0 +1,22 @@
+import torch
+from .GraphAttentionHeadModule import GraphAttentionHeadModule
+
+
+class MultiHeadGraphAttentionModule(torch.nn.Module):
+    def __init__(self, num_heads: int, in_features: int, out_features: int, alpha=0.2):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_out_features = out_features
+        self.heads = torch.nn.ModuleList(
+            [
+                GraphAttentionHeadModule(in_features, out_features, alpha=alpha)
+                for _ in range(num_heads)
+            ]
+        )
+
+    def forward(
+        self, node_features: torch.Tensor, adj_matrix: torch.Tensor
+    ) -> torch.Tensor:
+        outputs = [head(node_features, adj_matrix) for head in self.heads]
+        h_new_concat = torch.cat(outputs, dim=1)
+        return h_new_concat

homa/graph/modules/__init__.py

@@ -0,0 +1,2 @@
+from .GraphAttentionHeadModule import GraphAttentionHeadModule
+from .MultiHeadGraphAttentionModule import MultiHeadGraphAttentionModule

homa/loss/Loss.py

@@ -1,2 +1,5 @@
-class Loss:
+import torch
+
+
+class Loss(torch.nn.Module):
     pass

homa/rl/DQN.py

@@ -0,0 +1,2 @@
+class DQN:
+    pass

homa/rl/DRQN.py

@@ -0,0 +1,5 @@
+from .DQN import DQN
+
+
+class DRQN(DQN):
+    pass

homa/rl/DiversityIsAllYouNeed.py

@@ -0,0 +1,96 @@
+import torch
+from .diayn.Actor import Actor
+from .diayn.Critic import Critic
+from .diayn.Discriminator import Discriminator
+from .buffers import DiversityIsAllYouNeedBuffer, Buffer
+
+
+class DiversityIsAllYouNeed:
+    def __init__(
+        self,
+        state_dimension: int,
+        action_dimension: int,
+        hidden_dimension: int = 256,
+        num_skills: int = 10,
+        critic_decay: float = 0.0,
+        actor_decay: float = 0.0,
+        discriminator_decay: float = 0.0,
+        actor_lr: float = 0.0001,
+        critic_lr: float = 0.001,
+        discriminator_lr=0.001,
+        buffer_capacity: int = 1_000_000,
+        actor_epsilon: float = 1e-6,
+        gamma: float = 0.99,
+        min_std: float = -20.0,
+        max_std: float = 2.0,
+    ):
+        self.buffer: Buffer = DiversityIsAllYouNeedBuffer(capacity=buffer_capacity)
+        self.num_skills: int = num_skills
+        self.actor = Actor(
+            state_dimension=state_dimension,
+            action_dimension=action_dimension,
+            hidden_dimension=hidden_dimension,
+            num_skills=num_skills,
+            lr=actor_lr,
+            decay=actor_decay,
+            epsilon=actor_epsilon,
+            min_std=min_std,
+            max_std=max_std,
+        )
+        self.critic = Critic(
+            state_dimension=state_dimension,
+            hidden_dimension=hidden_dimension,
+            num_skills=num_skills,
+            lr=critic_lr,
+            decay=critic_decay,
+            gamma=gamma,
+        )
+        self.discriminator = Discriminator(
+            state_dimension=state_dimension,
+            hidden_dimension=hidden_dimension,
+            num_skills=num_skills,
+            lr=discriminator_lr,
+            decay=discriminator_decay,
+        )
+
+    def one_hot(self, indices, max_index) -> torch.Tensor:
+        one_hot = torch.zeros(indices.size(0), max_index)
+        one_hot.scatter_(1, indices.unsqueeze(1), 1)
+        return one_hot
+
+    def skill_index(self) -> torch.Tensor:
+        return torch.randint(0, self.num_skills, (1,))
+
+    def skill(self) -> torch.Tensor:
+        return self.one_hot(self.skill_index(), self.num_skills)
+
+    def advantages(
+        self,
+        states: torch.Tensor,
+        skills: torch.Tensor,
+        rewards: torch.Tensor,
+        terminations: torch.Tensor,
+        next_states: torch.Tensor,
+    ) -> torch.Tensor:
+        values = self.critic.values(states=states, skills=skills)
+        termination_mask = 1 - terminations
+        next_values = self.critic.values_(states=next_states, skills=skills)
+        update = self.gamma * next_values * termination_mask
+        return rewards + update - values
+
+    def train(self, skill: torch.Tensor):
+        data = self.buffer.all_tensor()
+        skill_indices = skill.repeat(data.states.size(0), 1).long()
+        skills_indices_one_hot = self.one_hot(skill_indices, self.num_skills)
+        self.discriminator.train(
+            states=data.states, skills_indices=skills_indices_one_hot
+        )
+        advantages = self.advantages(
+            states=data.states,
+            rewards=data.rewards,
+            terminations=data.terminations,
+            next_states=data.next_states,
+            skills=skills,
+        )
+        self.critic.train(advantages=advantages)
+        self.actor.train(advantages=advantages)

homa/rl/SoftActorCritic.py

@@ -0,0 +1,64 @@
+from .sac import SoftActor, SoftCritic
+from .buffers import SoftActorCriticBuffer
+
+
+class SoftActorCritic:
+    def __init__(
+        self,
+        state_dimension: int,
+        action_dimension: int,
+        hidden_dimension: int = 256,
+        buffer_capacity: int = 1_000_000,
+        batch_size: int = 256,
+        actor_lr: float = 0.0002,
+        critic_lr: float = 0.0003,
+        actor_decay: float = 0.0,
+        critic_decay: float = 0.0,
+        tau: float = 0.005,
+        alpha: float = 0.2,
+        gamma: float = 0.99,
+        min_std: float = -20,
+        max_std: float = 2,
+        warmup: int = 10_000,
+    ):
+        self.batch_size: int = batch_size
+        self.warmup: int = warmup
+
+        self.actor = SoftActor(
+            state_dimension=state_dimension,
+            action_dimension=action_dimension,
+            hidden_dimension=hidden_dimension,
+            lr=actor_lr,
+            weight_decay=actor_decay,
+            alpha=alpha,
+            min_std=min_std,
+            max_std=max_std,
+        )
+        self.critic = SoftCritic(
+            state_dimension=state_dimension,
+            action_dimension=action_dimension,
+            hidden_dimension=hidden_dimension,
+            lr=critic_lr,
+            weight_decay=critic_decay,
+            tau=tau,
+            gamma=gamma,
+            alpha=alpha,
+        )
+        self.buffer = SoftActorCriticBuffer(capacity=buffer_capacity)
+
+    def train(self):
+        # don't train before warmup
+        if self.buffer.size < self.warmup:
+            return
+
+        data = self.buffer.sample_torch(self.batch_size)
+        self.critic.train(
+            states=data.states,
+            actions=data.actions,
+            rewards=data.rewards,
+            terminations=data.terminations,
+            next_states=data.next_states,
+            actor=self.actor,
+        )
+        self.actor.train(states=data.states, critic_network=self.critic.network)
+        self.critic.update()

homa/rl/__init__.py

@@ -0,0 +1,4 @@
+from .DiversityIsAllYouNeed import DiversityIsAllYouNeed
+from .SoftActorCritic import SoftActorCritic
+from .DQN import DQN
+from .DRQN import DRQN

homa/rl/buffers/Buffer.py

@@ -0,0 +1,11 @@
+from .concerns import ResetsCollection, HasRecordAlternatives
+
+
+class Buffer(ResetsCollection, HasRecordAlternatives):
+    def __init__(self, capacity: int):
+        self.capacity: int = capacity
+        self.reset()
+
+    @property
+    def size(self):
+        return len(self.collection)

homa/rl/buffers/DiversityIsAllYouNeedBuffer.py

@@ -0,0 +1,50 @@
+import torch
+import numpy
+from types import SimpleNamespace
+from .Buffer import Buffer
+from .concerns import HasRecordAlternatives
+
+
+class DiversityIsAllYouNeedBuffer(Buffer, HasRecordAlternatives):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def all_tensor(self) -> SimpleNamespace:
+        return self.all(tensor=True)
+
+    def all(self, tensor: bool = False) -> SimpleNamespace:
+        states, actions, rewards, next_states, terminations, probabilities = zip(
+            *self.collection
+        )
+
+        if tensor:
+            states = torch.from_numpy(numpy.array(states))
+            actions = torch.from_numpy(numpy.array(actions))
+            rewards = torch.from_numpy(numpy.array(rewards))
+            next_states = torch.from_numpy(numpy.array(next_states))
+            terminations = torch.from_numpy(numpy.array(terminations))
+            probabilities = torch.from_numpy(numpy.array(probabilities))
+
+        return SimpleNamespace(
+            **{
+                "states": states,
+                "actions": actions,
+                "rewards": rewards,
+                "next_states": next_states,
+                "terminations": terminations,
+                "probabilities": probabilities,
+            }
+        )
+
+    def record(
+        self,
+        state: numpy.ndarray,
+        action: int,
+        reward: float,
+        next_state: numpy.ndarray,
+        termination: bool,
+        probability: numpy.ndarray,
+    ) -> None:
+        self.collection.append(
+            (state, action, reward, next_state, termination, probability)
+        )

homa/rl/buffers/ImageBuffer.py

@@ -0,0 +1,5 @@
+from .Buffer import Buffer
+
+
+class ImageBuffer(Buffer):
+    pass

homa/rl/buffers/SoftActorCriticBuffer.py

@@ -0,0 +1,56 @@
+import numpy
+import random
+import torch
+from types import SimpleNamespace
+from .Buffer import Buffer
+
+
+class SoftActorCriticBuffer(Buffer):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def record(
+        self,
+        state: numpy.ndarray,
+        action: int,
+        reward: float,
+        next_state: numpy.ndarray,
+        termination: float,
+        probability: numpy.ndarray,
+    ):
+        self.collection.append(
+            (state, action, reward, next_state, termination, probability)
+        )
+
+    def sample(self, k: int, as_tensor: bool = False):
+        batch = random.sample(self.collection, k)
+        states, actions, rewards, next_states, terminations, probabilities = zip(*batch)
+
+        states = numpy.array(states)
+        actions = numpy.array(actions)
+        rewards = numpy.array(rewards)
+        next_states = numpy.array(next_states)
+        terminations = numpy.array(terminations)
+        probabilities = numpy.array(probabilities)
+
+        if as_tensor:
+            states = torch.from_numpy(states).float()
+            actions = torch.from_numpy(actions).long()
+            rewards = torch.from_numpy(rewards).float()
+            next_states = torch.from_numpy(next_states).float()
+            terminations = torch.from_numpy(terminations).float()
+            probabilities = torch.from_numpy(probabilities).float()
+
+        return SimpleNamespace(
+            **{
+                "states": states,
+                "actions": actions,
+                "rewards": rewards,
+                "next_states": next_states,
+                "terminations": terminations,
+                "probabilities": probabilities,
+            }
+        )
+
+    def sample_torch(self, k: int):
+        return self.sample(k=k, as_tensor=True)

homa/rl/buffers/__init__.py

@@ -0,0 +1,4 @@
+from .SoftActorCriticBuffer import SoftActorCriticBuffer
+from .ImageBuffer import ImageBuffer
+from .DiversityIsAllYouNeedBuffer import DiversityIsAllYouNeedBuffer
+from .Buffer import Buffer

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)