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

fix: critic is no longer shared with actor network

Browse source
This commit is contained in:
cedric 2026-03-31 16:03:03 +00:00
parent 98a8f529e1
commit d119bf6408
1 changed files with 13 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

19
src/ppo.py
View file

@ -8,7 +8,9 @@ import jax.numpy as jnp
# Chose to use a class as it seemed the easiest way to integrate the CleanRL code style # Chose to use a class as it seemed the easiest way to integrate the CleanRL code style
# with our need to seperate concerns # with our need to seperate concerns
class PPO: class PPO:
def __init__(self, args, input_network, action_network, critic, message_passer=None): def __init__(
self, args, input_network, action_network, critic, critic_network, message_passer=None
):
self.args = args self.args = args
if not message_passer: if not message_passer:
@ -21,6 +23,7 @@ class PPO:
input_network_apply=input_network.apply, input_network_apply=input_network.apply,
action_network_apply=action_network.apply, action_network_apply=action_network.apply,
critic_apply=critic.apply, critic_apply=critic.apply,
critic_network_apply=critic_network.apply,
message_passer=message_passer, message_passer=message_passer,
), ),
has_aux=True, has_aux=True,
@ -83,24 +86,26 @@ that are now not in the same scope
""" """
@partial(jax.jit, static_argnums=(0, 1, 2, 3)) @partial(jax.jit, static_argnums=(0, 1, 2, 3, 4))
def get_action_and_value2( def get_action_and_value2(
input_apply, input_apply,
action_apply, action_apply,
message_passer, message_passer,
critic_apply, critic_apply,
critic_network_apply,
params: flax.core.FrozenDict, params: flax.core.FrozenDict,
x: jnp.ndarray, x: jnp.ndarray,
action: jnp.ndarray, action: jnp.ndarray,
): ):
hidden = input_apply(params["network_params"], x) hidden_network = input_apply(params["network_params"], x)
hidden = message_passer(hidden) hidden_critic = critic_network_apply(params["critic_network_params"], x)
mean, log_std = action_apply(params["actor_params"], hidden) hidden_network = message_passer(hidden_network)
mean, log_std = action_apply(params["actor_params"], hidden_network)
std = jnp.exp(log_std) std = jnp.exp(log_std)
logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1) logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1)
entropy = (0.5 + 0.5 * jnp.log(2 * jnp.pi) + log_std).sum(-1) entropy = (0.5 + 0.5 * jnp.log(2 * jnp.pi) + log_std).sum(-1)
value = critic_apply(params["critic_params"], hidden).squeeze(-1) value = critic_apply(params["critic_params"], hidden_critic).squeeze(-1)
return logprob, entropy, value return logprob, entropy, value
@ -117,12 +122,14 @@ def ppo_loss(
action_network_apply, action_network_apply,
message_passer, message_passer,
critic_apply, critic_apply,
critic_network_apply,
): ):
newlogprob, entropy, newvalue = get_action_and_value2( newlogprob, entropy, newvalue = get_action_and_value2(
input_network_apply, input_network_apply,
action_network_apply, action_network_apply,
message_passer, message_passer,
critic_apply, critic_apply,
critic_network_apply,
params, params,
x, x,
a, a,