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

fix: usage of field because flax wont allow mutable class object

Browse source
This commit is contained in:
cedric 2026-04-02 17:59:48 +00:00
parent 9879fea9e2
commit 95840b789b
1 changed files with 7 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

12
src/MLPs/mlps.py
View file

@ -1,4 +1,4 @@
from dataclasses import dataclass, fields from dataclasses import dataclass, fields, field
import flax import flax
import flax.linen as nn import flax.linen as nn
@ -11,8 +11,8 @@ from flax.linen.initializers import constant, orthogonal
# example usage: network = SemiGenericNetwork(layer_sizes=[256, 256], activation=nn.relu) # example usage: network = SemiGenericNetwork(layer_sizes=[256, 256], activation=nn.relu)
# semi generic so we can easily make a config for it in experiments # semi generic so we can easily make a config for it in experiments
class SemiGenericNetwork(nn.Module): class SemiGenericNetwork(nn.Module):
layer_sizes: Sequence[int] = [64, 64] # default 2 layers of 64 neurons layer_sizes: Sequence[int] = field(default_factory=lambda: [64, 64])
activation: Callable = nn.tanh # default tanh activation: Callable = nn.tanh
@nn.compact @nn.compact
def __call__(self, x): def __call__(self, x):
@ -29,11 +29,13 @@ class Critic(nn.Module):
class Actor(nn.Module): class Actor(nn.Module):
action_dim: Sequence[int] action_dim: int
@nn.compact @nn.compact
def __call__(self, x): def __call__(self, x):
return nn.Dense(self.action_dim, kernel_init=orthogonal(0.01), bias_init=constant(0.0))(x) mean = nn.Dense(self.action_dim, kernel_init=orthogonal(0.01), bias_init=constant(0.0))(x)
log_std = self.param("log_std", nn.initializers.zeros, (self.action_dim,))
return mean, log_std
@jax.tree_util.register_dataclass @jax.tree_util.register_dataclass