From 40c7e32344a727e364e2830549f3120826d7353a Mon Sep 17 00:00:00 2001 From: cedric Date: Thu, 2 Apr 2026 18:11:49 +0000 Subject: [PATCH] feat: more generic naming for critic, since messager will use the same --- src/MLPs/mlps.py | 2 +- src/train.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/MLPs/mlps.py b/src/MLPs/mlps.py index 6c587d3..1bae0c4 100644 --- a/src/MLPs/mlps.py +++ b/src/MLPs/mlps.py @@ -21,7 +21,7 @@ class SemiGenericNetwork(nn.Module): return x -class Critic(nn.Module): +class OneDenseLayerMLP(nn.Module): @nn.compact def __call__(self, x): return nn.Dense(1, kernel_init=orthogonal(1), bias_init=constant(0.0))(x) diff --git a/src/train.py b/src/train.py index 89b85d5..7528f26 100644 --- a/src/train.py +++ b/src/train.py @@ -19,7 +19,7 @@ from torch.utils.tensorboard import SummaryWriter from brittle_star_project.dataclasses import PPOArgs from brittle_star_project.dataclasses.EpisodeStatistics import EpisodeStatistics from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper -from MLPs.mlps import SemiGenericNetwork, Actor, Critic, AgentParams, Storage +from MLPs.mlps import SemiGenericNetwork, Actor, OneDenseLayerMLP, AgentParams, Storage from ppo import PPO @@ -119,7 +119,8 @@ def train(args: PPOArgs): network = SemiGenericNetwork() critic_network = SemiGenericNetwork() actor = Actor(action_dim=env.single_action_space.shape[0]) # continuous actions for MJX - critic = Critic() + critic = OneDenseLayerMLP() + # messager = OneDenseLayerMLP() sample_obs = jnp.concatenate( [