streaming-deep-rl 0.2.0__tar.gz → 0.2.2__tar.gz

{streaming_deep_rl-0.2.0 → streaming_deep_rl-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: streaming-deep-rl
-Version: 0.2.0
+Version: 0.2.2
 Summary: Streaming Deep Reinforcement Learning
 Project-URL: Homepage, https://pypi.org/project/streaming-deep-rl/
 Project-URL: Repository, https://codeberg.org/lucidrains/streaming-deep-rl

{streaming_deep_rl-0.2.0 → streaming_deep_rl-0.2.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "streaming-deep-rl"
-version = "0.2.0"
+version = "0.2.2"
 description = "Streaming Deep Reinforcement Learning"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{streaming_deep_rl-0.2.0 → streaming_deep_rl-0.2.2}/streaming_deep_rl/streaming_deep_rl.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 import math
 from collections import deque
+from contextlib import nullcontext
 from typing import Callable, NamedTuple
 
 import torch
@@ -25,12 +26,21 @@ from streaming_deep_rl.buffer_dict import BufferDict
 def exists(v):
     return v is not None
 
+def identity(t):
+    return t
+
 def divisible_by(num, den):
     return (num % den) == 0
 
 def default(v, d):
     return v if exists(v) else d
 
+def no_grad_detach(fn):
+    def inner(t):
+        with torch.no_grad():
+            return fn(t).detach()
+    return inner
+
 def cast_tensor(t):
     return tensor(t) if not is_tensor(t) else t
 
@@ -104,7 +114,8 @@ class SelfPredictRepr(Module):
         dim_embed,
         num_discrete_actions = 0,
         num_continuous_actions = 0,
-        dim_predict = None
+        dim_predict = None,
+        target_embed_from_ema = True # whether the target embed is from an EMA, or from the model itself (traditional vs sigreg from lejepa)
     ):
         super().__init__()
         assert num_discrete_actions > 0 or num_continuous_actions > 0
@@ -113,6 +124,8 @@ class SelfPredictRepr(Module):
 
         self.to_action_embed = Embed(dim_embed, num_discrete = num_discrete_actions, num_continuous = num_continuous_actions)
 
+        self.target_embed_from_ema = target_embed_from_ema
+
         self.to_projection = nn.Sequential(
             nn.RMSNorm(dim_embed),
             nn.Linear(dim_embed, dim_predict, bias = False)
@@ -130,8 +143,11 @@ class SelfPredictRepr(Module):
     ):
         projected = self.to_projection(state_embed)
 
-        with torch.no_grad():
-            next_projected = self.to_projection(next_state_embed.detach())
+        # maybe no
+
+        decorator = no_grad_detach if self.target_embed_from_ema else identity
+
+        next_projected = decorator(self.to_projection)(next_state_embed)
 
         action_embed = self.to_action_embed(actions)
 
@@ -314,7 +330,8 @@ class StreamingACLambda(Module):
         actor_use_ema = False,
         actor_ema_beta = 0.7,
         critic_ema_beta = 0.7,
-        self_predict_repr = False,
+        actor_self_predict_repr = False,
+        critic_self_predict_repr = False,
         entropy_weight = 0.01,
         init_sparsity = 0.9,
         dim_critic = 128,
@@ -384,18 +401,29 @@ class StreamingACLambda(Module):
         # self-predictive representation
         # Schwarzer et al. https://arxiv.org/abs/2007.05929
 
-        self.self_predict_repr = self_predict_repr
+        self.actor_self_predict_repr = actor_self_predict_repr
+        self.critic_self_predict_repr = critic_self_predict_repr
 
-        if self_predict_repr:
+        if actor_self_predict_repr:
 
             actor_use_ema |= True
 
-            self.actor_self_predict_repr = SelfPredictRepr(
+            self.actor_spr = SelfPredictRepr(
                 dim_actor,
                 num_discrete_actions = num_discrete_actions,
                 num_continuous_actions = num_continuous_actions,
             )
 
+        if critic_self_predict_repr:
+
+            use_critic_ema |= True
+
+            self.critic_spr = SelfPredictRepr(
+                dim_critic,
+                num_discrete_actions = num_discrete_actions,
+                num_continuous_actions = num_continuous_actions,
+            )
+
         # actor
 
         self.actor = actor
@@ -652,7 +680,7 @@ class StreamingACLambda(Module):
             # get the SPR loss
 
             actor_spr_grad = None
-            if self.self_predict_repr:
+            if self.actor_self_predict_repr:
                 actor_embed = self.actor(oldest_state)
 
                 if self.actor_use_ema:
@@ -660,9 +688,9 @@ class StreamingACLambda(Module):
                 else:
                     actor_next_embed = self.actor(last_next_state)
 
-                spr_loss = self.actor_self_predict_repr(actor_embed, oldest_action, actor_next_embed)
+                spr_loss = self.actor_spr(actor_embed, oldest_action, actor_next_embed)
 
-                spr_params = list(self.actor_self_predict_repr.parameters())
+                spr_params = list(self.actor_spr.parameters())
                 spr_grads_tup = torch.autograd.grad(spr_loss, actor_params + spr_params, retain_graph = True, allow_unused = True)
                 spr_grads_tup = tuple(default(g, torch.zeros_like(p)) for g, p in zip(spr_grads_tup, actor_params + spr_params))
 
@@ -673,6 +701,30 @@ class StreamingACLambda(Module):
                 for param, grad in zip(spr_params, spr_only_grad):
                     param.data.add_(grad, alpha = -self.actor_lr)
 
+            # critic SPR loss
+
+            critic_spr_grad = None
+            if self.critic_self_predict_repr:
+                critic_embed = self.critic(oldest_state)
+
+                if self.use_critic_ema:
+                    critic_next_embed = self.critic_ema.ema_model[0](last_next_state)
+                else:
+                    critic_next_embed = self.critic(last_next_state)
+
+                critic_spr_loss = self.critic_spr(critic_embed, oldest_action, critic_next_embed)
+
+                critic_spr_params = list(self.critic_spr.parameters())
+                critic_spr_grads_tup = torch.autograd.grad(critic_spr_loss, critic_params + critic_spr_params, retain_graph = True, allow_unused = True)
+                critic_spr_grads_tup = tuple(default(g, torch.zeros_like(p)) for g, p in zip(critic_spr_grads_tup, critic_params + critic_spr_params))
+
+                critic_spr_grad = {name: grad for (name, _), grad in zip(self.critic_full.named_parameters(), critic_spr_grads_tup[:len(critic_params)])}
+                critic_spr_only_grad = critic_spr_grads_tup[len(critic_params):]
+
+                # Update SPR networks via SGD
+                for param, grad in zip(critic_spr_params, critic_spr_only_grad):
+                    param.data.add_(grad, alpha = -self.critic_lr)
+
             # entropy gradient decoupled from td_error
 
             entropy_grads_tup = torch.autograd.grad(entropy, actor_params, retain_graph = True)
@@ -820,7 +872,7 @@ class StreamingACLambda(Module):
             self.cautious_wd,
             self.wd_towards_init,
             None,
-            None
+            critic_spr_grad
         )
 
         if self.actor_use_ema: