From c5a8dd2b3aeabd86cd214b0e086958fb50ea72b0 Mon Sep 17 00:00:00 2001
From: JibrilExe <cedric00ik@gmail.com>
Date: Sun, 12 Apr 2026 12:54:48 +0200
Subject: [PATCH] feat: expanded used var doc, and added distance to target log

---
 configs/hpc/debug.yaml                        |  2 +-
 .../used_variables.md                         | 16 +++++++++++++-
 .../environment/env_config.py                 |  2 +-
 .../trainers/PPOTrainer.py                    | 21 ++++++++++++++++---
 4 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/configs/hpc/debug.yaml b/configs/hpc/debug.yaml
index c387dc1..e2a9f9d 100644
--- a/configs/hpc/debug.yaml
+++ b/configs/hpc/debug.yaml
@@ -4,7 +4,7 @@ seed: 42
 track: true
 wandb_project_name: "Let's-find-that-bug"
 wandb_entity: "SEL3-2026-Groep-4"
- 
+run_dir: "/data/gent/465/vsc46589"
 num_envs: 32
 num_steps: 32
 total_timesteps: 102400
diff --git a/experiments/debug-experiment-10042026/used_variables.md b/experiments/debug-experiment-10042026/used_variables.md
index 0dcba35..e49c9e4 100644
--- a/experiments/debug-experiment-10042026/used_variables.md
+++ b/experiments/debug-experiment-10042026/used_variables.md
@@ -1,6 +1,6 @@
 ## Default envconfig
 task: Task = Task.DIRECTED_LOCOMOTION
-simulation_time: float = 5.0
+simulation_time: float = 500.0
 num_physics_steps_per_control_step: int = 10
 time_scale: int = 2
 camera_ids: list[int] = field(default_factory=lambda: [0, 1])
@@ -53,6 +53,8 @@ use_torque_control: bool = False
 
 ## MLPs:
 ### Sensor & Feature_extractor:
+Both with 3 layers of 300 neurons per layer.
+
 class GenericDenseLayersWithActivation(nn.Module):
     layer_sizes: Sequence[int] = field(default_factory=lambda: [64, 64])
     activation: Callable = nn.tanh
@@ -80,3 +82,15 @@ class OneDenseLayerMLP(nn.Module):
         return nn.Dense(1, kernel_init=orthogonal(1), bias_init=constant(0.0))(x)
 
 ### Observations:
+_ALLOWED_OBS_KEYS = {
+    "joint_position",
+    "joint_velocity",
+    "joint_actuator_force",
+    "actuator_force",
+    "disk_position",
+    "disk_rotation",
+    "disk_linear_velocity",
+    "disk_angular_velocity",
+    "unit_xy_direction_to_target",
+    "xy_distance_to_target",
+}
\ No newline at end of file
diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py
index 78083e9..c2e8067 100644
--- a/src/brittle_star_project/environment/env_config.py
+++ b/src/brittle_star_project/environment/env_config.py
@@ -31,7 +31,7 @@ class EnvConfig:
 
     task: Task = Task.DIRECTED_LOCOMOTION
 
-    simulation_time: float = 5.0
+    simulation_time: float = 500.0
     num_physics_steps_per_control_step: int = 10
     time_scale: int = 2
 
diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py
index 2dc02f9..6997911 100644
--- a/src/brittle_star_project/trainers/PPOTrainer.py
+++ b/src/brittle_star_project/trainers/PPOTrainer.py
@@ -38,6 +38,12 @@ _ALLOWED_OBS_KEYS = {
 }
 # TODO: clip scaled reward?
 
+@jax.jit
+def _get_xy_distance_to_target(obs_dict: dict) -> jnp.ndarray:
+    """Extract xy_distance_to_target for all environments."""
+    # obs_dict is a dict of arrays with leading batch dimension (num_envs, ...)
+    return obs_dict["xy_distance_to_target"].squeeze(-1)  # shape: (num_envs,)
+
 @jax.jit
 def _clip_action(action: jnp.ndarray, low: jnp.ndarray, high: jnp.ndarray) -> jnp.ndarray:
     return jnp.clip(action, low, high)
@@ -143,7 +149,7 @@ 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 *= 100
+    reward *= 1000
     terminated = next_env_state.terminated
     truncated = next_env_state.truncated
     done = terminated | truncated
@@ -396,7 +402,9 @@ class PPOTrainer:
         start_time,
         iteration_time_start,
         training_measurements,
-        storage
+        storage,
+        next_obs,
+        xy_distance
     ):
         data = jax.device_get({
             'rewards': storage.rewards[0],      # (num_steps,)
@@ -425,6 +433,9 @@ class PPOTrainer:
 
             "rollout/env0/action_mean": float(np.mean(data['actions'])),
             "rollout/env0/raw_action_mean": float(np.mean(data['raw_actions'])),
+
+            "charts/env0_xy_distance_to_target": float(xy_distance[0]),
+            "charts/env1_xy_distance_to_target": float(xy_distance[1]),
         }
 
         metrics = {
@@ -568,6 +579,8 @@ class PPOTrainer:
                 env_state, next_obs, next_done, iteration=iteration
             )
 
+            xy_distance = _get_xy_distance_to_target(env_state.observations)
+
             global_step += self.args.num_steps * self.args.num_envs
             self._log(
                 global_step,
@@ -575,7 +588,9 @@ class PPOTrainer:
                 start_time,
                 iteration_time_start,
                 training_measurements,
-                storage
+                storage,
+                next_obs,
+                xy_distance
             )
 
             sps = int(global_step / (time.time() - start_time))