tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
Archived
1
Fork 0
Code Issues Releases Activity

feat: generic network

Browse source
This commit is contained in:
cedric 2026-04-02 17:38:46 +00:00
parent 2263c97fe8
commit 15ae86796f
5 changed files with 9 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

31
src/MLPs/centralized.py
View file

@ -1,31 +0,0 @@
from typing import Sequence
import flax.linen as nn
import jax.numpy as jnp
from flax.linen.initializers import constant, orthogonal
class Network(nn.Module):
hidden_size: int = 256
@nn.compact
def __call__(self, x):
# x shape: (batch, obs_dim)
x = nn.Dense(self.hidden_size, kernel_init=orthogonal(jnp.sqrt(2)))(x)
x = nn.tanh(x)
x = nn.Dense(self.hidden_size, kernel_init=orthogonal(jnp.sqrt(2)))(x)
x = nn.tanh(x)
return x
class Critic(nn.Module):
@nn.compact
def __call__(self, x):
return nn.Dense(1, kernel_init=orthogonal(1), bias_init=constant(0.0))(x)
class Actor(nn.Module):
action_dim: Sequence[int]
@nn.compact
def __call__(self, x):
return nn.Dense(self.action_dim, kernel_init=orthogonal(0.01), bias_init=constant(0.0))(x)