From 0a9bf2e0a55f8d02c35e451cc2812a002c5e6d88 Mon Sep 17 00:00:00 2001
From: JibrilExe <cedric00ik@gmail.com>
Date: Tue, 14 Apr 2026 20:30:21 +0200
Subject: [PATCH] fix: dont store clipped action, and upgraded reward scale

---
 src/brittle_star_project/ppo.py                 | 15 ++-------------
 src/brittle_star_project/trainers/PPOTrainer.py |  5 +++--
 2 files changed, 5 insertions(+), 15 deletions(-)

diff --git a/src/brittle_star_project/ppo.py b/src/brittle_star_project/ppo.py
index ffe2064..0156b82 100644
--- a/src/brittle_star_project/ppo.py
+++ b/src/brittle_star_project/ppo.py
@@ -143,19 +143,8 @@ def ppo_loss(
     pg_loss1 = -mb_advantages * ratio
     pg_loss2 = -mb_advantages * jnp.clip(ratio, 1 - args.clip_coef, 1 + args.clip_coef)
     pg_loss = jnp.maximum(pg_loss1, pg_loss2).mean()
-    old_value = mb_returns - mb_advantages
-
-    v_clipped = old_value + jnp.clip(
-        newvalue - old_value,
-        -args.clip_coef,
-        args.clip_coef
-    )
-
-    v_loss_unclipped = (newvalue - mb_returns) ** 2
-    v_loss_clipped = (v_clipped - mb_returns) ** 2
-
-    v_loss = 0.5 * jnp.maximum(v_loss_unclipped, v_loss_clipped).mean()
-
+    
+    v_loss = 0.5 * ((newvalue - mb_returns) ** 2).mean()
     entropy_loss = entropy.mean()
     loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef
     return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl))
diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py
index ec19e26..3dcf6db 100644
--- a/src/brittle_star_project/trainers/PPOTrainer.py
+++ b/src/brittle_star_project/trainers/PPOTrainer.py
@@ -134,7 +134,7 @@ def _step_once(
 
     storage = Storage(
         obs=obs,
-        actions=clipped_action,
+        actions=raw_action,
         raw_actions=raw_action,
         logprobs=logprob,
         dones=done,
@@ -152,7 +152,8 @@ def _step_env_wrapped(episode_stats, env_state, action, env_step_fn):
     next_env_state = env_step_fn(env_state, action)
 
     reward = next_env_state.reward
-    reward *= 1000
+    reward *= 5000
+    reward = jnp.clip(reward, -1, 1)
     terminated = next_env_state.terminated
     truncated = next_env_state.truncated
     done = terminated | truncated