Port TRPO, PPO, GAIL

scratch/etienne/trpo/core/value.py

@@ -0,0 +1,48 @@
+import torch
+import torch.nn as nn
+from torch.distributions import Normal
+from torch.distributions.kl import kl_divergence
+
+class Value(nn.Module):
+    
+    def __init__(self):
+        super().__init__()
+        self.nn = nn.Sequential(
+            nn.LazyLinear(50),
+            nn.Tanh(),
+            nn.LazyLinear(50),
+            nn.Tanh(),
+            nn.LazyLinear(1),
+        )
+    
+    def forward(self, states):
+        return self.nn(states).squeeze(-1)
+
+class SetValue(Value):
+
+    def forward(self, states):
+        batch_size = states.shape[:-2]
+        states = torch.cat((states[..., :1, [0, 1]], states[..., :, [2, 5]]), axis=-2).reshape(*batch_size, -1)
+        return super().forward(states)
+
+class DeepSetValue(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.elem = nn.Sequential(
+            nn.LazyLinear(10),
+            nn.Tanh(),
+            nn.LazyLinear(10),
+            nn.Tanh(),
+            nn.LazyLinear(10),
+        )
+        self.glob = nn.Sequential(
+            nn.LazyLinear(10),
+            nn.Tanh(),
+            nn.LazyLinear(10),
+            nn.Tanh(),
+            nn.LazyLinear(1),
+        )
+    
+    def forward(self, states):
+        return self.glob(self.elem(states).sum(-2)).squeeze(-1)