locoformer 0.0.5__py3-none-any.whl → 0.0.15__py3-none-any.whl

locoformer/locoformer.py

@@ -1,12 +1,25 @@
 from __future__ import annotations
 from functools import partial
 
+from pathlib import Path
+from contextlib import contextmanager
+from collections import namedtuple
+
+import numpy as np
+from numpy import ndarray
+from numpy.lib.format import open_memmap
+
+from beartype import beartype
+from beartype.door import is_bearable
+
 import torch
-from torch import cat, stack, is_tensor
+from torch import nn, cat, stack, arange, Tensor, tensor, is_tensor, from_numpy
 import torch.nn.functional as F
-from torch.nn import Module, ModuleList, Linear, RMSNorm, Identity
+from torch.nn import Module, ModuleList, Linear, RMSNorm, Identity, Sequential
 from torch.utils._pytree import tree_map
+from torch.utils.data import Dataset, DataLoader
 
+import einx
 from einops import rearrange, einsum
 from einops.layers.torch import Rearrange
 
@@ -24,23 +37,39 @@ def exists(v):
 def default(v, d):
     return v if exists(v) else d
 
+def first(arr):
+    return arr[0]
+
 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 combine_kv_cache(cache1, cache2):
-    combined_cache = []
+def pad_at_dim(
+    t,
+    pad: tuple[int, int],
+    dim = -1,
+    value = 0.
+):
+    if pad == (0, 0):
+        return t
 
-    for layer_cache1, layer_cache2 in zip(cache1, cache2):
-        next_cache = cat((layer_cache1, layer_cache2), dim = -2)
-        combined_cache.append(next_cache)
+    dims_from_right = (- dim - 1) if dim < 0 else (t.ndim - dim - 1)
+    zeros = ((0, 0) * dims_from_right)
+    return F.pad(t, (*zeros, *pad), value = value)
 
-    return combined_cache
+def calc_entropy(logits):
+    prob = logits.softmax(dim = -1)
+    return -(prob * log(prob)).sum(dim = -1)
 
 # generalized advantage estimate
 
@@ -48,7 +77,7 @@ def combine_kv_cache(cache1, cache2):
 def calc_gae(
     rewards,
     values,
-    masks,
+    masks = None,
     gamma = 0.99,
     lam = 0.95,
     use_accelerated = None
@@ -59,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
 
@@ -110,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
@@ -142,15 +174,229 @@ def create_sliding_mask(
     create_kwargs = dict(device = device) if exists(device) else dict()
     return create_block_mask(sliding_mask, B = None, H = None, Q_LEN = seq_len, KV_LEN = kv_seq_len, _compile = True, **create_kwargs)
 
+# data
+
+def collate_var_time(data):
+
+    datum = first(data)
+    keys = datum.keys()
+
+    all_tensors = zip(*[datum.values() for datum in data])
+
+    collated_values = []
+
+    for key, tensors in zip(keys, all_tensors):
+
+        # the episode lens have zero dimension - think of a cleaner way to handle this later
+
+        if key != '_lens':
+
+            times = [t.shape[0] for t in tensors]
+            max_time = max(times)
+            tensors = [pad_at_dim(t, (0, max_time - t.shape[0]), dim = 0) for t in tensors]
+
+        collated_values.append(stack(tensors))
+
+    return dict(zip(keys, collated_values))
+
+class ReplayDataset(Dataset):
+    def __init__(
+        self,
+        folder: str | Path,
+        fields: tuple[str, ...] | None = None
+    ):
+        if isinstance(folder, str):
+            folder = Path(folder)
+
+        episode_lens = folder / 'episode_lens.npy'
+        self.episode_lens = open_memmap(str(episode_lens), mode = 'r')
+
+        # get indices of non-zero lengthed episodes
+
+        nonzero_episodes = self.episode_lens > 0
+        self.indices = np.arange(self.episode_lens.shape[-1])[nonzero_episodes]
+
+        # get all data files
+
+        filepaths = [*folder.glob('*.data.npy')]
+        assert len(filepaths) > 0
+
+        fieldname_to_filepath = {path.name.split('.')[0]: path for path in filepaths}
+
+        fieldnames_from_files = set(fieldname_to_filepath.keys())
+
+        fields = default(fields, fieldnames_from_files)
+
+        self.memmaps = dict()
+
+        for field in fields:
+            assert field in fieldnames_from_files, f'invalid field {field} - must be one of {fieldnames_from_files}'
+
+            path = fieldname_to_filepath[field]
+
+            self.memmaps[field] = open_memmap(str(path), mode = 'r')
+
+    def __len__(self):
+        return len(self.indices)
+
+    def __getitem__(self, idx):
+        episode_index = self.indices[idx]
+
+        episode_len = self.episode_lens[episode_index]
+
+        data = {field: from_numpy(memmap[episode_index, :episode_len].copy()) for field, memmap in self.memmaps.items()}
+
+        data['_lens'] = tensor(episode_len)
+
+        return data
+
+class ReplayBuffer:
+
+    @beartype
+    def __init__(
+        self,
+        folder: str | Path,
+        max_episodes: int,
+        max_timesteps: int,
+        fields: dict[
+            str,
+            str | tuple[str, int | tuple[int, ...]]
+        ]
+    ):
+
+        # folder for data
+
+        if not isinstance(folder, Path):
+            folder = Path(folder)
+            folder.mkdir(exist_ok = True)
+
+        self.folder = folder
+        assert folder.is_dir()
+
+        # keeping track of episode length
+
+        episode_lens = folder / 'episode_lens.npy'
+
+        self.episode_index = 0
+        self.timestep_index = 0
+
+        self.max_episodes = max_episodes
+        self.max_timesteps= max_timesteps
+
+        self.episode_lens = open_memmap(str(episode_lens), mode = 'w+', dtype = np.int32, shape = (max_episodes,))
+
+        # create the memmap for individual data tracks
+
+        self.shapes = dict()
+        self.dtypes = dict()
+        self.memmaps = dict()
+        self.fieldnames = set(fields.keys())
+
+        for field_name, field_info in fields.items():
+
+            # some flexibility
+
+            field_info = (field_info, ()) if isinstance(field_info, str) else field_info
+
+            dtype_str, shape = field_info
+            assert dtype_str in {'int', 'float', 'bool'}
+
+            dtype = dict(int = np.int32, float = np.float32, bool = np.bool_)[dtype_str]
+
+            # memmap file
+
+            filepath = folder / f'{field_name}.data.npy'
+            memmap = open_memmap(str(filepath), mode = 'w+', dtype = dtype, shape = (max_episodes, max_timesteps, *shape))
+
+            self.memmaps[field_name] = memmap
+            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
+
+    def flush(self):
+        self.episode_lens[self.episode_index] = self.timestep_index
+
+        for memmap in self.memmaps.values():
+            memmap.flush()
+
+        self.episode_lens.flush()
+
+    @contextmanager
+    def one_episode(self):
+
+        yield
+
+        self.flush()
+        self.advance_episode()
+
+    @beartype
+    def store_datapoint(
+        self,
+        episode_index: int,
+        timestep_index: int,
+        name: str,
+        datapoint: Tensor | ndarray
+    ):
+        assert 0 <= episode_index < self.max_episodes
+        assert 0 <= timestep_index < self.max_timesteps
+
+        if is_tensor(datapoint):
+            datapoint = datapoint.detach().cpu().numpy()
+
+        assert name in self.fieldnames, f'invalid field name {name} - must be one of {self.fieldnames}'
+
+        assert datapoint.shape == self.shapes[name], f'invalid shape {datapoint.shape} - shape must be {self.shapes[name]}'
+
+        self.memmaps[name][self.episode_index, self.timestep_index] = datapoint
+
+    def store(
+        self,
+        **data
+    ):
+        assert is_bearable(data, dict[str, Tensor | ndarray])
+
+        assert not self.timestep_index >= self.max_timesteps, 'you exceeded the `max_timesteps` set on the replay buffer'
+
+        for name, datapoint in data.items():
+
+            self.store_datapoint(self.episode_index, self.timestep_index, name, datapoint)
+
+        self.timestep_index += 1
+
+        return self.memory_namedtuple(**data)
+
+    def dataset(self) -> Dataset:
+        self.flush()
+
+        return ReplayDataset(self.folder)
+
+    def dataloader(self, batch_size, **kwargs) -> DataLoader:
+        self.flush()
+
+        return DataLoader(self.dataset(), batch_size = batch_size, collate_fn = collate_var_time, **kwargs)
+
 # transformer-xl with ppo
 
 class Attention(Module):
     def __init__(
         self,
         dim,
+        window_size,
         dim_head = 64,
         heads = 8,
-        pre_rmsnorm = True
+        pre_rmsnorm = True,
+        fixed_window_size = False,
+        accept_value_residual = False
     ):
         super().__init__()
         self.scale = dim_head ** -0.5
@@ -167,20 +413,54 @@ class Attention(Module):
         self.to_kv = LinearNoBias(dim, dim_inner * 2)
         self.to_out = LinearNoBias(dim_inner, dim)
 
+        self.to_v_gates = Sequential(
+            LinearNoBias(dim, heads),
+            Rearrange('b n h -> b h n 1'),
+            nn.Sigmoid()
+        )
+
+        # value residual
+
+        self.accept_value_residual = accept_value_residual
+
+        if accept_value_residual:
+            self.to_value_residual_mix = Sequential(
+                LinearNoBias(dim, heads),
+                Rearrange('b n h -> b h n 1'),
+                nn.Sigmoid()                
+            )
+
+        # fixed window size
+
+        self.fixed_window_size = fixed_window_size
+        self.window_size = window_size
+
     def forward(
         self,
         tokens,
+        value_residual = None,
         kv_cache = None,
-        return_kv_cache = False
+        return_kv_cache = False,
     ):
+        seq_len = tokens.shape[-2]
+
+        device = tokens.device
+
         tokens = self.norm(tokens)
 
         q, k, v = (self.to_q(tokens), *self.to_kv(tokens).chunk(2, dim = -1))
 
         q, k, v = map(self.split_heads, (q, k, v))
 
+        orig_v = v
+
         q = q * self.scale
 
+        if exists(value_residual):
+            assert self.accept_value_residual
+            mix = self.to_value_residual_mix(tokens)
+            v = v.lerp(value_residual, mix)
+
         if exists(kv_cache):
             ck, cv = kv_cache
             k = cat((ck, k), dim = -2)
@@ -195,7 +475,13 @@ class Attention(Module):
 
         i, j = sim.shape[-2:]
 
-        causal_mask = torch.ones((i, j), dtype = torch.bool, device = sim.device).triu(j - i + 1)
+        if self.fixed_window_size:
+            i_seq = arange(i, device = device)
+            j_seq = arange(j, device = device) - (j - i)
+            dist = einx.subtract('i, j -> i j', i_seq, j_seq)
+            causal_mask = (dist < 0) | (dist > self.window_size)
+        else:
+            causal_mask = torch.ones((i, j), dtype = torch.bool, device = sim.device).triu(j - i + 1)
 
         sim = sim.masked_fill(causal_mask, -torch.finfo(sim.dtype).max)
 
@@ -203,6 +489,8 @@ class Attention(Module):
 
         out = einsum(attn, v, 'b h i j, b h j d -> b h i d')
 
+        out = out * self.to_v_gates(tokens)
+
         out = self.merge_heads(out)
 
         out = self.to_out(out)
@@ -210,7 +498,7 @@ class Attention(Module):
         if not return_kv_cache:
             return out
 
-        return out, next_kv_cache
+        return out, (next_kv_cache, orig_v)
 
 class FeedForward(Module):
     def __init__(
@@ -244,17 +532,21 @@ class TransformerXL(Module):
         self,
         dim,
         depth,
+        window_size,
         dim_head = 64,
         heads = 8,
         expansion_factor = 4.,
-        final_norm = True
+        final_norm = True,
+        fixed_window_size = False,
     ):
         super().__init__()
 
         layers = ModuleList([])
 
-        for _ in range(depth):
-            attn = Attention(dim = dim, dim_head = dim_head, heads = heads)
+        for i in range(depth):
+            is_first = i == 0
+
+            attn = Attention(dim = dim, dim_head = dim_head, heads = heads, fixed_window_size = fixed_window_size, window_size = window_size, accept_value_residual = not is_first)
 
             ff = FeedForward(dim = dim, expansion_factor = expansion_factor)
 
@@ -265,6 +557,11 @@ class TransformerXL(Module):
         self.layers = layers
         self.norm = RMSNorm(dim) if final_norm else Identity()
 
+        # fixed window size
+
+        self.fixed_window_size = fixed_window_size
+        self.window_size = window_size
+
     def forward(
         self,
         x,
@@ -275,22 +572,28 @@ class TransformerXL(Module):
         cache = default(cache, (None,) * len(self.layers))
 
         next_kv_caches = []
+        value_residual = None
 
         for (attn, ff), kv_cache in zip(self.layers, cache):
 
-            attn_out, next_kv_cache = attn(x, kv_cache = kv_cache, return_kv_cache = True)
-
-            next_kv_caches.append(next_kv_cache)
+            attn_out, (next_kv_cache, values) = attn(x, value_residual = value_residual, kv_cache = kv_cache, return_kv_cache = True)
 
             x = attn_out + x
             x = ff(x) + x
 
+            next_kv_caches.append(next_kv_cache)
+            value_residual = default(value_residual, values)
+
         embed = self.norm(x)
 
         if not return_kv_cache:
             return embed
 
-        return embed, stack(next_kv_caches)
+        next_kv_cache = stack(next_kv_caches)
+
+        next_kv_cache = next_kv_cache[..., -self.window_size:, :]
+
+        return embed, next_kv_cache
 
 # class
 
@@ -300,7 +603,13 @@ 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
     ):
         super().__init__()
 
@@ -314,28 +623,138 @@ class Locoformer(Module):
 
         self.value_network = value_network
 
+        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
+
     @property
     def device(self):
         return next(self.parameters()).device
 
+    def actor_parameters(self):
+        return self.unembedder.parameters()
+
+    def critic_parameters(self):
+        if not exists(self.value_network):
+            return []
+
+        return self.value_network.parameters()
+
+    def ppo(
+        self,
+        state,
+        action,
+        old_action_log_prob,
+        reward,
+        old_value,
+        mask,
+        actor_optim,
+        critic_optim
+    ):
+
+        (action_logits, value), _ = self.forward(state, 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()
+
+        returns = calc_gae(reward, old_value, lam = self.gae_lam, gamma = self.discount_factor)
+        advantage = returns - old_value
+
+        actor_loss = -torch.min(ratio * advantage, ratio.clamp(1. - eps_clip, 1. + eps_clip) * advantage)
+
+        actor_loss = actor_loss - self.ppo_entropy_weight * entropy
+
+        mean_actor_loss = actor_loss[mask].mean()
+        mean_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
+
+        mean_critic_loss = critic_loss[mask].mean()
+        mean_critic_loss.backward()
+
+        # optimizer update
+
+        actor_optim.step()
+        actor_optim.zero_grad()
+
+        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):
+            state, _ =  env.reset(*args, **kwargs)
+
+            if isinstance(state, ndarray):
+                state = from_numpy(state)
+
+            return state, _
+
+        def wrapped_step(action, *args, **kwargs):
+            out = env.step(action.item(), *args, **kwargs)
+
+            def transform_output(el):
+                if isinstance(el, ndarray):
+                    return from_numpy(el)
+                elif isinstance(el, (int, bool, float)):
+                    return tensor(el)
+                else:
+                    return el
+
+            return tree_map(transform_output, out)
+
+        return wrapped_reset, wrapped_step
+
     def get_stateful_forward(
         self,
-        segment_size,
         initial_states: Tensor | None = None,
         inference_mode = False,
         has_batch_dim = False,
+        has_time_dim = False,
         **kwargs
     ):
+        window_size = self.window_size
+
         cache = None
 
-        def stateful_forward(state: Tensor, override_kwargs: dict = dict()):
+        def stateful_forward(state: Tensor, **override_kwargs):
             nonlocal cache
 
-            # handle no batch, for easier time rolling out against envs
+            # handle no batch or time, for easier time rolling out against envs
 
             if not has_batch_dim:
                 state = rearrange(state, '... -> 1 ...')
 
+            if not has_time_dim:
+                state = rearrange(state, '... d -> ... 1 d')
+
             # forwards
 
             out, cache = self.forward(state, cache = cache, **{**kwargs, **override_kwargs})
@@ -344,10 +763,13 @@ class Locoformer(Module):
 
             cache_len = cache.shape[-2]
 
-            if divisible_by(cache_len, segment_size * 2):
-                cache = cache[..., -segment_size:, :]
+            if self.fixed_window_size or divisible_by(cache_len, window_size * 2):
+                cache = cache[..., -window_size:, :]
 
-            # maybe remove batch
+            # maybe remove batch or time
+
+            if not has_time_dim:
+                out = tree_map_tensor(out, lambda t: rearrange(t, '... 1 d -> ... d'))
 
             if not has_batch_dim:
                 out = tree_map_tensor(out, lambda t: rearrange(t, '1 ... -> ...'))
@@ -364,7 +786,7 @@ class Locoformer(Module):
 
         initial_logits = []
 
-        for state_segments in initial_states.split(segment_size, dim = -1):
+        for state_segments in initial_states.split(self.window_size, dim = -1):
 
             logits = stateful_forward(state_segments, return_values = False)
             initial_logits.append(logits)
@@ -381,6 +803,8 @@ class Locoformer(Module):
         return_values = False
     ):
 
+        state = state.to(self.device)
+
         tokens = self.embedder(state)
 
         embed, kv_cache = self.transformer(tokens, cache = cache, return_kv_cache = True)

{locoformer-0.0.5.dist-info → locoformer-0.0.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: locoformer
-Version: 0.0.5
+Version: 0.0.15
 Summary: LocoFormer
 Project-URL: Homepage, https://pypi.org/project/locoformer/
 Project-URL: Repository, https://github.com/lucidrains/locoformer
@@ -35,6 +35,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
 Requires-Dist: assoc-scan
+Requires-Dist: beartype
 Requires-Dist: einops>=0.8.0
 Requires-Dist: einx>=0.3.0
 Requires-Dist: rotary-embedding-torch
@@ -53,7 +54,7 @@ Description-Content-Type: text/markdown
 
 [LocoFormer - Generalist Locomotion via Long-Context Adaptation](https://generalist-locomotion.github.io/)
 
-The gist is they trained a simple Transformer-XL in simulation on robots with many different bodies (cross-embodiment). When transferring to the real-world, they noticed the robot now gains the ability to adapt to insults. The XL memories span across multiple trials, which allowed the robot to learn in-context adaptation.
+The gist is they trained a simple Transformer-XL in simulation on robots with many different bodies (cross-embodiment) with extreme domain randomization. When transferring to the real-world, they noticed the robot now gains the ability to adapt to insults. The XL memories span across multiple trials, which allowed the robot to learn in-context adaptation.
 
 ## Sponsors
 

locoformer-0.0.15.dist-info/RECORD

@@ -0,0 +1,6 @@
+locoformer/__init__.py,sha256=XctsMGEZSR4mVl75fhds_1BtS5qGFiiItTDV7CmCt_I,45
+locoformer/locoformer.py,sha256=1jPK41G4HB1PEPtlusQxcrne489E-3QKXAULZ20FEZM,22740
+locoformer-0.0.15.dist-info/METADATA,sha256=IHtK7NvVQewYQ0GBB7v1KG90_H2Jakxir0MakUIA-jU,3218
+locoformer-0.0.15.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+locoformer-0.0.15.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+locoformer-0.0.15.dist-info/RECORD,,

locoformer-0.0.5.dist-info/RECORD

@@ -1,6 +0,0 @@
-locoformer/__init__.py,sha256=XctsMGEZSR4mVl75fhds_1BtS5qGFiiItTDV7CmCt_I,45
-locoformer/locoformer.py,sha256=Yoh3hrj2E_91YLoYRa73wGzjdIiMdcd5ofNjkiVlogI,10570
-locoformer-0.0.5.dist-info/METADATA,sha256=oe6HfOwWKQvusiJl1ukmNFcrGRhdDZ6NcKZi3upv-SY,3159
-locoformer-0.0.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-locoformer-0.0.5.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-locoformer-0.0.5.dist-info/RECORD,,