feat: first approach to action clipping

configs/hpc/debug.yaml

@@ -8,6 +8,7 @@ wandb_entity: "SEL3-2026-Groep-4"
 num_envs: 32
 num_steps: 32
 total_timesteps: 102400
-
+num_arms: 2
+num_segments_per_arm: 1
 cuda: true
 

experiments/debug-experiment-10042026/used_variables.md

@@ -46,7 +46,7 @@ wall_height: float = 1.5
 wall_thickness: float = 0.1
 
 ## Morphology:
-num_arms: int = 5
+num_arms: int = 2
 num_segments_per_arm: int = 4
 use_p_control: bool = True
 use_torque_control: bool = False

scripts/hpc/train.pbs

@@ -65,8 +65,8 @@ fi
 
 # TODO Once experiments get serious, change the config
 python scripts/train.py \
-    --env-config-path configs/hpc/smoke_test.yaml \
+    --env-config-path configs/hpc/debug.yaml \
-    --hyperparameter-config-path configs/hpc/smoke_test.yaml \
+    --hyperparameter-config-path configs/hpc/debug.yaml \
     --run-dir "$SCRATCH_RUNDIR"
 
 echo ">>> Staging out results to $DATA_RUNDIR..."

scripts/train.py

@@ -69,18 +69,10 @@ if __name__ == "__main__":
 
     env = make_env(args.env_config_path, args.num_envs)
     raw_env = env.raw
-    print(
+    logger.log(
-        "\n\n\n Observation space \n",
+        {"run_dir": run_dir}
-        raw_env.observation_space,
-        "Action space \n",
-        raw_env.action_space,
-    )
-    print(
-        "\n\n\n Observation space \n",
-        raw_env.observation_space,
-        "Action space \n",
-        raw_env.action_space,
     )
+
     torch.backends.cudnn.deterministic = args.torch_deterministic
 
     ppo_trainer = PPOTrainer(args, env, run_dir, run_name)

src/brittle_star_project/trainers/PPOTrainer.py

@@ -24,6 +24,9 @@ from brittle_star_project.MLPs.mlps import (
 )
 from brittle_star_project.ppo import PPO
 
+@jax.jit
+def _clip_action(action: jnp.ndarray, low: jnp.ndarray, high: jnp.ndarray) -> jnp.ndarray:
+    return jnp.clip(action, low, high)
 
 @jax.jit
 def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, learning_rate):
@@ -47,6 +50,8 @@ def _get_action_and_value_noise(
     agent_state: TrainState,
     next_obs: jnp.ndarray,
     key: jax.random.PRNGKey,
+    action_low,
+    action_high
 ):
     hidden = sensor.apply(agent_state.params["sensor_params"], next_obs)
     hidden_critic = feature_extractor.apply(
@@ -59,9 +64,14 @@ def _get_action_and_value_noise(
     noise = jax.random.normal(subkey, shape=mean.shape)
     std = jnp.exp(log_std)
     action = mean + noise * std
-    logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1)
+    clipped_action = _clip_action(
+        action,
+        action_low,
+        action_high
+    )
+    logprob = -0.5 * (((clipped_action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1)
     value = critic.apply(agent_state.params["critic_params"], hidden_critic)
-    return action, logprob, value.squeeze(-1), key
+    return clipped_action, logprob, value.squeeze(-1), key
 
 
 # removed jit: used in _rollout_jit, so will be compiled with _rollout_jit
@@ -73,10 +83,12 @@ def _step_once(
     feature_extractor: GenericDenseLayersWithActivation,
     actor: Actor,
     critic: OneDenseLayerMLP,
+    action_low,
+    action_high
 ):
     agent_state, episode_stats, obs, done, key, env_state = carry
     action, logprob, value, key = _get_action_and_value_noise(
-        sensor, feature_extractor, actor, critic, agent_state, obs, key
+        sensor, feature_extractor, actor, critic, agent_state, obs, key, action_low, action_high
     )
 
     episode_stats, env_state, (next_obs, reward, next_done) = env_step_fn(
@@ -140,6 +152,8 @@ def _rollout_jit(
     feature_extractor: GenericDenseLayersWithActivation,
     actor: Actor,
     critic: OneDenseLayerMLP,
+    action_low,
+    action_high
 ):
     (agent_state, episode_stats, next_obs, next_done, key, env_state), storage = jax.lax.scan(
         partial(
@@ -149,6 +163,8 @@ def _rollout_jit(
             actor=actor,
             critic=critic,
             env_step_fn=step_env_fn,
+            action_low=action_low,
+            action_high=action_high
         ),
         (agent_state, episode_stats, next_obs, next_done, key, env_state),
         (),
@@ -215,6 +231,9 @@ class PPOTrainer:
         self.actor.apply = jax.jit(self.actor.apply)
         self.critic.apply = jax.jit(self.critic.apply)
 
+        action_low = jnp.asarray(self.env.single_action_space.low, dtype=jnp.float32)
+        action_high = jnp.asarray(self.env.single_action_space.high, dtype=jnp.float32)
+
         self._rollout_jit = jax.jit(
             partial(
                 _rollout_jit,
@@ -224,6 +243,8 @@ class PPOTrainer:
                 feature_extractor=self.feature_extractor,
                 actor=self.actor,
                 critic=self.critic,
+                action_low=action_low,
+                action_high=action_high
             )
         )
         self._compute_gae_jit = jax.jit(