From 95840b789b01a7e4ef6ef7484d8a67e99c4fb164 Mon Sep 17 00:00:00 2001 From: cedric Date: Thu, 2 Apr 2026 17:59:48 +0000 Subject: [PATCH] fix: usage of field because flax wont allow mutable class object --- src/MLPs/mlps.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/MLPs/mlps.py b/src/MLPs/mlps.py index e4b5969..f49ac5c 100644 --- a/src/MLPs/mlps.py +++ b/src/MLPs/mlps.py @@ -1,4 +1,4 @@ -from dataclasses import dataclass, fields +from dataclasses import dataclass, fields, field import flax 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) # semi generic so we can easily make a config for it in experiments class SemiGenericNetwork(nn.Module): - layer_sizes: Sequence[int] = [64, 64] # default 2 layers of 64 neurons - activation: Callable = nn.tanh # default tanh + layer_sizes: Sequence[int] = field(default_factory=lambda: [64, 64]) + activation: Callable = nn.tanh @nn.compact def __call__(self, x): @@ -29,11 +29,13 @@ class Critic(nn.Module): class Actor(nn.Module): - action_dim: Sequence[int] + action_dim: int @nn.compact 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