locoformer 0.0.11__tar.gz → 0.0.17__tar.gz

{locoformer-0.0.11 → locoformer-0.0.17}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: locoformer
-Version: 0.0.11
+Version: 0.0.17
 Summary: LocoFormer
 Project-URL: Homepage, https://pypi.org/project/locoformer/
 Project-URL: Repository, https://github.com/lucidrains/locoformer

{locoformer-0.0.11 → locoformer-0.0.17}/locoformer/locoformer.py

@@ -3,6 +3,7 @@ from functools import partial
 
 from pathlib import Path
 from contextlib import contextmanager
+from collections import namedtuple
 
 import numpy as np
 from numpy import ndarray
@@ -17,6 +18,7 @@ import torch.nn.functional as F
 from torch.nn import Module, ModuleList, Linear, RMSNorm, Identity, Sequential
 from torch.utils._pytree import tree_map
 from torch.utils.data import Dataset, DataLoader
+from torch.optim import Optimizer
 
 import einx
 from einops import rearrange, einsum
@@ -26,6 +28,8 @@ from rotary_embedding_torch import RotaryEmbedding
 
 from assoc_scan import AssocScan
 
+# constants
+
 LinearNoBias = partial(Linear, bias = False)
 
 # helper functions
@@ -42,12 +46,14 @@ def first(arr):
 def divisible_by(num, den):
     return (num % den) == 0
 
+# tensor helpers
+
+def log(t, eps = 1e-20):
+    return t.clamp_min(eps).log()
+
 def tree_map_tensor(x, fn):
     return tree_map(lambda t: t if not is_tensor(t) else fn(t), x)
 
-def detach_all(x):
-    return tree_map_tensor(x, lambda t: t.detach())
-
 def pad_at_dim(
     t,
     pad: tuple[int, int],
@@ -61,13 +67,17 @@ def pad_at_dim(
     zeros = ((0, 0) * dims_from_right)
     return F.pad(t, (*zeros, *pad), value = value)
 
+def calc_entropy(logits):
+    prob = logits.softmax(dim = -1)
+    return -(prob * log(prob)).sum(dim = -1)
+
 # generalized advantage estimate
 
 @torch.no_grad()
 def calc_gae(
     rewards,
     values,
-    masks,
+    masks = None,
     gamma = 0.99,
     lam = 0.95,
     use_accelerated = None
@@ -78,6 +88,9 @@ def calc_gae(
     values = F.pad(values, (0, 1), value = 0.)
     values, values_next = values[..., :-1], values[..., 1:]
 
+    if not exists(masks):
+        masks = torch.ones_like(values)
+
     delta = rewards + gamma * values_next * masks - values
     gates = gamma * lam * masks
 
@@ -87,7 +100,7 @@ def calc_gae(
 
     returns = gae + values
 
-    return returns
+    return gae, returns
 
 # transformer-xl mask w/ flex attn
 
@@ -129,8 +142,8 @@ def create_xl_mask(
         # handle intra-episodic attention if needed
 
         if exists(episode_ids):
-            q_episode = episodes[b, q + offset]
-            k_episode = episodes[b, k]
+            q_episode = episode_ids[b, q + offset]
+            k_episode = episode_ids[b, k]
 
             intra_episode_mask = q_episode == k_episode
             mask = mask & intra_episode_mask
@@ -231,7 +244,7 @@ class ReplayDataset(Dataset):
 
         episode_len = self.episode_lens[episode_index]
 
-        data = {field: torch.from_numpy(memmap[episode_index, :episode_len]) for field, memmap in self.memmaps.items()}
+        data = {field: from_numpy(memmap[episode_index, :episode_len].copy()) for field, memmap in self.memmaps.items()}
 
         data['_lens'] = tensor(episode_len)
 
@@ -299,6 +312,13 @@ class ReplayBuffer:
             self.shapes[field_name] = shape
             self.dtypes[field_name] = dtype
 
+        self.memory_namedtuple = namedtuple('Memory', list(fields.keys()))
+
+    def reset_(self):
+        self.episode_lens[:] = 0
+        self.episode_index = 0
+        self.timestep_index = 0
+
     def advance_episode(self):
         self.episode_index = (self.episode_index + 1) % self.max_episodes
         self.timestep_index = 0
@@ -353,15 +373,17 @@ class ReplayBuffer:
 
         self.timestep_index += 1
 
+        return self.memory_namedtuple(**data)
+
     def dataset(self) -> Dataset:
         self.flush()
 
         return ReplayDataset(self.folder)
 
-    def dataloader(self, **kwargs) -> DataLoader:
+    def dataloader(self, batch_size, **kwargs) -> DataLoader:
         self.flush()
 
-        return DataLoader(self.dataset(), collate_fn = collate_var_time, **kwargs)
+        return DataLoader(self.dataset(), batch_size = batch_size, collate_fn = collate_var_time, **kwargs)
 
 # transformer-xl with ppo
 
@@ -421,7 +443,6 @@ class Attention(Module):
         return_kv_cache = False,
     ):
         seq_len = tokens.shape[-2]
-        assert seq_len <= self.window_size
 
         device = tokens.device
 
@@ -582,7 +603,14 @@ class Locoformer(Module):
         embedder: Module,
         unembedder: Module,
         transformer: dict | TransformerXL,
-        value_network: Module | None = None
+        value_network: Module | None = None,
+        discount_factor = 0.999,
+        gae_lam = 0.95,
+        ppo_eps_clip = 0.2,
+        ppo_entropy_weight = 0.01,
+        ppo_value_clip = 0.4,
+        value_loss_weight = 0.5,
+        calc_gae_kwargs: dict = dict()
     ):
         super().__init__()
 
@@ -599,6 +627,21 @@ class Locoformer(Module):
         self.fixed_window_size = transformer.fixed_window_size
         self.window_size = transformer.window_size
 
+        # ppo related
+
+        self.discount_factor = discount_factor
+        self.gae_lam = gae_lam
+        self.ppo_eps_clip = ppo_eps_clip
+        self.ppo_entropy_weight = ppo_entropy_weight
+        self.ppo_value_clip = ppo_value_clip
+        self.value_loss_weight = value_loss_weight
+
+        self.calc_gae_kwargs = calc_gae_kwargs
+
+        # loss related
+
+        self.register_buffer('zero', tensor(0.), persistent = False)
+
     @property
     def device(self):
         return next(self.parameters()).device
@@ -612,6 +655,101 @@ class Locoformer(Module):
 
         return self.value_network.parameters()
 
+    def ppo(
+        self,
+        state,
+        action,
+        old_action_log_prob,
+        reward,
+        old_value,
+        mask,
+        actor_optim: Optimizer | None = None,
+        critic_optim: Optimizer | None = None
+    ):
+        window_size = self.window_size
+        total_learnable_tokens = mask.sum().item()
+
+        windowed_tensors = [
+            t.split(window_size, dim = 1) for t in
+            (
+                state,
+                action,
+                old_action_log_prob,
+                reward,
+                old_value,
+                mask
+            )
+        ]
+
+        mean_actor_loss = self.zero.clone()
+        mean_critic_loss = self.zero.clone()
+
+        # learn across windows
+
+        cache = None
+
+        for (
+            state,
+            action,
+            old_action_log_prob,
+            reward,
+            old_value,
+            mask
+        ) in zip(*windowed_tensors):
+
+            (action_logits, value), cache = self.forward(state, cache = cache, detach_cache = True, return_values = True)
+            entropy = calc_entropy(action_logits)
+
+            action = rearrange(action, 'b t -> b t 1')
+            log_prob = action_logits.gather(-1, action)
+            log_prob = rearrange(log_prob, 'b t 1 -> b t')
+
+            # update actor, classic clipped surrogate loss
+
+            eps_clip = self.ppo_eps_clip
+            ratio = (log_prob - old_action_log_prob).exp()
+
+            advantage, returns = calc_gae(reward, old_value, lam = self.gae_lam, gamma = self.discount_factor, **self.calc_gae_kwargs)
+
+            actor_loss = -torch.min(ratio * advantage, ratio.clamp(1. - eps_clip, 1. + eps_clip) * advantage)
+
+            actor_loss = actor_loss - self.ppo_entropy_weight * entropy
+
+            windowed_actor_loss = actor_loss[mask].sum() / total_learnable_tokens
+            windowed_actor_loss.backward(retain_graph = True)
+
+            # update critic
+
+            value_loss = F.mse_loss(returns, value, reduction = 'none')
+
+            value_clip = self.ppo_value_clip
+            clipped_value = old_value + (value - old_value).clamp(-value_clip, value_clip)
+            clipped_value_loss = F.mse_loss(returns, clipped_value, reduction = 'none')
+
+            critic_loss = torch.maximum(value_loss, clipped_value_loss) * self.value_loss_weight
+
+            windowed_critic_loss = critic_loss[mask].sum() / total_learnable_tokens
+            windowed_critic_loss.backward(retain_graph = True)
+
+            # accumulate
+
+            mean_actor_loss.add_(windowed_actor_loss)
+            mean_critic_loss.add_(windowed_critic_loss)
+
+        # optimizer update
+
+        if exists(actor_optim):
+            actor_optim.step()
+            actor_optim.zero_grad()
+
+        if exists(critic_optim):
+            critic_optim.step()
+            critic_optim.zero_grad()
+
+        # return losses for logging
+
+        return mean_actor_loss.detach(), mean_critic_loss.detach()
+
     def wrap_env_functions(self, env):
 
         def wrapped_reset(*args, **kwargs):
@@ -723,7 +861,7 @@ class Locoformer(Module):
         # maybe detach cache
 
         if detach_cache:
-            kv_cache = detach_all(kv_cache)
+            kv_cache = kv_cache.detach()
 
         # handle returning of values
 

{locoformer-0.0.11 → locoformer-0.0.17}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "locoformer"
-version = "0.0.11"
+version = "0.0.17"
 description = "LocoFormer"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{locoformer-0.0.11 → locoformer-0.0.17}/train.py

@@ -169,7 +169,7 @@ for i in range(NUM_BATCHES):
         prime = prime.to(model.device)
         out = prime
 
-        stateful_forward, logits = model.get_stateful_forward(has_batch_dim = False, initial_states = prime, inference_mode = True)
+        stateful_forward, logits = model.get_stateful_forward(has_batch_dim = False, has_time_dim = True, initial_states = prime, inference_mode = True)
 
         # sample
 

{locoformer-0.0.11 → locoformer-0.0.17}/train_gym.py

@@ -3,7 +3,7 @@
 #     "accelerate",
 #     "fire",
 #     "gymnasium[box2d]>=1.0.0",
-#     "locoformer",
+#     "locoformer>=0.0.12",
 #     "moviepy",
 #     "tqdm"
 # ]
@@ -13,17 +13,19 @@ from fire import Fire
 from shutil import rmtree
 from tqdm import tqdm
 from collections import deque
+from types import SimpleNamespace
 
 from accelerate import Accelerator
 
 import gymnasium as gym
 
 import torch
-from torch import from_numpy, randint, tensor, stack
+from torch import from_numpy, randint, tensor, stack, arange
 import torch.nn.functional as F
 from torch.utils.data import TensorDataset, DataLoader
 from torch.optim import Adam
 
+import einx
 from einops import rearrange
 
 from locoformer.locoformer import Locoformer, ReplayBuffer
@@ -47,26 +49,69 @@ def gumbel_sample(logits, temperature = 1., eps = 1e-6):
     noise = gumbel_noise(logits)
     return ((logits / max(temperature, eps)) + noise).argmax(dim = -1)
 
+# learn
+
+def learn(
+    model,
+    actor_optim,
+    critic_optim,
+    accelerator,
+    replay,
+    batch_size = 16,
+    epochs = 2,
+):
+    device = accelerator.device
+
+    dl = replay.dataloader(batch_size = batch_size, shuffle = True)
+    model, dl, actor_optim, critic_optim = accelerator.prepare(model, dl, actor_optim, critic_optim)
+
+    for _ in range(epochs):
+        for data in dl:
+
+            data = SimpleNamespace(**data)
+
+            seq_len = data.state.shape[1]
+
+            value_mask = einx.less('j, i -> i j', arange(seq_len, device = device), data._lens)
+            value = torch.where(value_mask, data.value, 0.)
+
+            actor_loss, critic_loss = model.ppo(
+                state = data.state,
+                action = data.action,
+                old_action_log_prob = data.action_log_prob,
+                reward = data.reward,
+                old_value = value,
+                mask = data.learnable,
+                actor_optim = actor_optim,
+                critic_optim = critic_optim
+            )
+
+            accelerator.print(f'actor: {actor_loss.item():.3f} | critic: {critic_loss.item():.3f}')
+
 # main function
 
 def main(
     env_name = 'LunarLander-v3',
     num_episodes = 50_000,
     max_timesteps = 500,
-    num_timestep_before_learn = 5000,
+    num_episodes_before_learn = 32,
     clear_video = True,
     video_folder = 'recordings',
     record_every_episode = 250,
+    learning_rate = 8e-4,
     discount_factor = 0.99,
-    learning_rate = 1e-4,
+    betas = (0.9, 0.99),
+    gae_lam = 0.95,
+    ppo_eps_clip = 0.2,
+    ppo_entropy_weight = .01,
     batch_size = 16,
     epochs = 2
 ):
 
     # accelerate
 
-    accelerate = Accelerator()
-    device = accelerate.device
+    accelerator = Accelerator()
+    device = accelerator.device
 
     # environment
 
@@ -91,14 +136,15 @@ def main(
     replay = ReplayBuffer(
         'replay',
         num_episodes,
-        max_timesteps,
+        max_timesteps + 1, # one extra node for bootstrap node - not relevant for locoformer, but for completeness
         fields = dict(
             state = ('float', (dim_state,)),
             action = 'int',
             action_log_prob = 'float',
             reward = 'float',
             value = 'float',
-            done = 'bool'
+            done = 'bool',
+            learnable = 'bool'
         )
     )
 
@@ -114,11 +160,18 @@ def main(
             heads = 4,
             depth = 4,
             window_size = 16
+        ),
+        discount_factor = discount_factor,
+        gae_lam = gae_lam,
+        ppo_eps_clip = ppo_eps_clip,
+        ppo_entropy_weight = ppo_entropy_weight,
+        calc_gae_kwargs = dict(
+            use_accelerated = False
         )
     ).to(device)
 
-    optim_actor = Adam([*locoformer.transformer.parameters(), *locoformer.actor_parameters()], lr = learning_rate)
-    optim_critic = Adam([*locoformer.transformer.parameters(), *locoformer.critic_parameters()], lr = learning_rate)
+    optim_actor = Adam([*locoformer.transformer.parameters(), *locoformer.actor_parameters()], lr = learning_rate, betas = betas)
+    optim_critic = Adam([*locoformer.transformer.parameters(), *locoformer.critic_parameters()], lr = learning_rate, betas = betas)
 
     timesteps_learn = 0
 
@@ -129,7 +182,8 @@ def main(
 
     # loop
 
-    for _ in tqdm(range(num_episodes)):
+    for episodes_index in tqdm(range(num_episodes)):
+
         state, *_ = env_reset()
 
         timestep = 0
@@ -158,27 +212,29 @@ def main(
                 action_log_prob = action_logits.gather(-1, rearrange(action, '-> 1'))
                 action_log_prob = rearrange(action_log_prob, '1 ->')
 
-                replay.store(
+                memory = replay.store(
                     state = state,
                     action = action,
                     action_log_prob = action_log_prob,
                     reward = reward,
                     value = value,
-                    done = done
+                    done = done,
+                    learnable = tensor(True)
                 )
 
-                # increment counters
+                # handle bootstrap value, which is a non-learnable timestep added with the next value for GAE
+                # only if terminated signal not detected
 
-                timestep += 1
-                timesteps_learn += 1
+                if not terminated:
+                    _, next_value = stateful_forward(next_state, return_values = True)
+
+                    memory._replace(value = next_value, learnable = False)
 
-                # learn if hit the number of learn timesteps
+                    replay.store(**memory._asdict())
 
-                if timesteps_learn >= num_timestep_before_learn:
-                    # todo - carry out learning
+                # increment counters
 
-                    timesteps_learn = 0
-                    memories.clear()
+                timestep += 1
 
                 # break if done or exceed max timestep
 
@@ -187,6 +243,19 @@ def main(
 
                 state = next_state
 
+            # learn if hit the number of learn timesteps
+
+            if divisible_by(episodes_index + 1, num_episodes_before_learn):
+
+                learn(
+                    locoformer,
+                    optim_actor,
+                    optim_critic,
+                    accelerator,
+                    replay,
+                    batch_size,
+                    epochs,
+                )
 # main
 
 if __name__ == '__main__':