fix: dont store clipped action, and upgraded reward scale

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))

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