fix: robot direction to target

docs/design/input_action_spaces.md

@@ -16,7 +16,8 @@ Global inputs, always broadcasted to all nodes:
   $$
   tilt = sqrt(roll^2 + pitch^2)
   $$
-- Goal vector: Instead of just a scalar distance, the goal is represented as an angle/direction to the target.
+- Goal vector: A 2D unit vector representing the *egocentric* direction to the target. A value of $[1.0, 0.0]$
+  indicates that the target is directly in front of the agent (angle 0).
 
 Local inputs, routed directly to specific nodes:
 
@@ -73,20 +74,21 @@ When designing the state space, we must ask: *Could a human operator perform thi
   **NOTE:** We later dropped the "distance to vector", switching to only a direction as the input. Our reasoning is
   the agent should always move towards the goal (it should not learn to stop at the goal), which allows for this
   simplification that decreases the model input size.
+  
+  The environment provides a raw `unit_xy_direction_to_target` (global), which we transform into a calculated
+  `robot_direction_to_target` (egocentric) before passing it to the MLPs. This vector consists of the X and Y
+  direction, where a value of $[1.0, 0.0]$ (mapping to an angle of $0$) means the robot is facing directly towards the
+  target.
 - Contact sensors: Segment contact detects external ground interaction and is biologically vital for timing gait
   transitions.
-- **Zero-Centered Rescaling ($[-1, 1]$):** Using a zero-centered range is standard best practice for continuous control
+- Zero-Centered Rescaling ($[-1, 1]$): Using a zero-centered range is standard best practice for continuous control
   tasks. It provides several mathematical and physical advantages:
-  - **Improved Gradient Flow:** Neural networks optimize faster when inputs are zero-centered. If all inputs were
-    positive (e.g., $[0, 1]$), the gradients during backpropagation would be forced to the same sign, causing
-    inefficient "zig-zag" weight updates.
-  - **Meaningful "Neutral" State:** In robotics, $0.0$ naturally represents a resting state (zero velocity, centered
+  - Improved Gradient Flow: Neural networks optimize faster when inputs are zero-centered. If all inputs were positive
+    (e.g., $[0, 1]$), the gradients during backpropagation would be forced to the same sign, causing inefficient
+    "zig-zag" weight updates.
+  - Meaningful Neutral State: In robotics, $0.0$ naturally represents a resting state (zero velocity, centered
     position, no force). In a $[-1, 1]$ system, this physical rest maps to a neutral $0.0$ signal in the network.
-  - **Robustness to Amputation:** In this project, amputated limbs are padded with $0.0$. In a $[-1, 1]$ system, this
-    correctly communicates a "neutral/dead" signal. In a $[0, 1]$ system, $0.0$ would represent the absolute minimum
-    physical limit, causing the network to misinterpret missing limbs as being at their extreme limits.
-
-
+    This also correctly communicaties a "neutral/dead" signal for amputated limbs that are padded with $0.0$ values.
 
 Specifically, we do not include some available inputs:
 
@@ -120,8 +122,7 @@ This is what the filtered input vectors look like in MuJoCo, with $J$ joints and
 - `joint_velocity`: shape=(J,), dtype=float64
 - `joint_actuator_force`: shape=(J,), dtype=float64
 - `segment_contact`: shape=(S,), dtype=float64
-- `unit_xy_direction_to_target`: shape=(2,), dtype=float64
-- `xy_distance_to_target`: shape=(1,), dtype=float64
+- `robot_direction_to_target`: shape=(2,), dtype=float64, egocentric
 - `disk_z_tilt`: shape=(1,), dtype=float64, derived from `disk_rotation`
 
 This brings the entire input space down to $3J + S + 4$ float64's, compared to $4J + S + 15$ float64's for the
@@ -159,6 +160,5 @@ disk_angular_velocity: shape=(3,), dtype=float64, size=3
 tendon_position: shape=(0,), dtype=float64, size=0
 tendon_velocity: shape=(0,), dtype=float64, size=0
 segment_contact: shape=(6,), dtype=float64, size=6
-unit_xy_direction_to_target: shape=(2,), dtype=float64, size=2
 xy_distance_to_target: shape=(1,), dtype=float64, size=1
 ```

src/brittle_star_project/environment/env_config.py

@@ -70,21 +70,21 @@ class ObservationBoundsConfig:
 
     # Based on max. ctrlrange (0.78539816339744828) in XML, but empirical testing went slightly over
     joint_position: list[float] = field(default_factory=lambda: [-0.8, 0.8])
-    # Empirical testing showed max. 3.22, adding buffer to be safe
+    # Empirical testing showed max. 3.22, adding buffer to be safe. Consider higher values "fast".
     joint_velocity: list[float] = field(default_factory=lambda: [-5.0, 5.0])
     # Based on max. forceRange in XML, verified with empirical testing
     joint_actuator_force: list[float] = field(default_factory=lambda: [-3.75, 3.75])
     # Based on intuition and reasoning
     segment_contact: list[float] = field(default_factory=lambda: [0.0, 1.0])
-    unit_xy_direction_to_target: list[float] = field(default_factory=lambda: [-1.0, 1.0])
+    robot_direction_to_target: list[float] = field(default_factory=lambda: [-1.0, 1.0])
     disk_z_tilt: list[float] = field(default_factory=lambda: [0.0, 3.141592653589793])
 
     def to_bounds_dict(self) -> dict[str, tuple[float, float]]:
         return {
+            "disk_z_tilt": tuple(self.disk_z_tilt),
+            "joint_actuator_force": tuple(self.joint_actuator_force),
             "joint_position": tuple(self.joint_position),
             "joint_velocity": tuple(self.joint_velocity),
-            "joint_actuator_force": tuple(self.joint_actuator_force),
+            "robot_direction_to_target": tuple(self.robot_direction_to_target),
             "segment_contact": tuple(self.segment_contact),
-            "unit_xy_direction_to_target": tuple(self.unit_xy_direction_to_target),
-            "disk_z_tilt": tuple(self.disk_z_tilt),
         }

src/brittle_star_project/environment/obs_processing.py

@@ -26,6 +26,15 @@ def create_obs_processor(
         if "disk_rotation" in new_obs:
             rot = new_obs["disk_rotation"]
             new_obs["disk_z_tilt"] = jnp.sqrt(jnp.pow(rot[0], 2) + jnp.pow(rot[1], 2))
+
+            if "unit_xy_direction_to_target" in new_obs:
+                yaw = rot[2]
+                unit_x, unit_y = new_obs["unit_xy_direction_to_target"]
+                cos_yaw, sin_yaw = jnp.cos(yaw), jnp.sin(yaw)
+                new_x = unit_x * cos_yaw + unit_y * sin_yaw
+                new_y = -unit_x * sin_yaw + unit_y * cos_yaw
+                new_obs["robot_direction_to_target"] = jnp.stack([new_x, new_y])
+
         return new_obs
 
     def _normalize_features(obs: dict) -> dict:
@@ -61,8 +70,8 @@ def create_obs_processor(
             "joint_actuator_force",
             "joint_position",
             "joint_velocity",
+            "robot_direction_to_target",
             "segment_contact",
-            "unit_xy_direction_to_target",
         ]
         values = []
         for key in ordered_keys:

tests/test_target_direction.py

@@ -0,0 +1,78 @@
+import jax.numpy as jnp
+
+from brittle_star_project.configs.main_config import BrittleStarConfig
+from brittle_star_project.environment.env_types import Backend
+from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
+from brittle_star_project.environment.obs_processing import create_obs_processor
+
+
+def test_raw_environment_returns_allocentric_direction():
+    """
+    Verifies that the raw environment returns a GLOBAL (allocentric)
+    direction to the target. If the robot rotates in place,
+    the global vector to the target should remain identical.
+    """
+    env = BrittleStarJaxEnvWrapper.default(num_envs=1, backend=Backend.MJX)
+    env_state = env.reset(seed=42)
+    raw_obs_1 = env_state.observations["unit_xy_direction_to_target"]
+
+    ninety_deg_z_quat = jnp.array([0.7071068, 0.0, 0.0, 0.7071068])
+    new_qpos = env_state.mjx_data.qpos.at[..., 3:7].set(ninety_deg_z_quat)
+    new_data = env_state.mjx_data.replace(qpos=new_qpos)
+    rotated_env_state = env_state.replace(mjx_data=new_data)
+
+    zero_action = jnp.zeros(env.single_action_space.shape)
+    if len(raw_obs_1.shape) > 1:
+        zero_action = jnp.expand_dims(zero_action, 0)
+    final_env_state = env.step(rotated_env_state, zero_action)
+    raw_obs_2 = final_env_state.observations["unit_xy_direction_to_target"]
+
+    # If the vector is allocentric, it should not change when the robot spins.
+    assert jnp.sum(jnp.abs(raw_obs_1 - raw_obs_2)) < 1e-4, (
+        f"The raw environment observation changed when the robot rotated! "
+        f"This means it is already egocentric. "
+        f"Obs 1: {raw_obs_1}, Obs 2: {raw_obs_2}"
+    )
+
+
+def test_processor_converts_to_egocentric_direction():
+    """
+    Verifies that the obs_processor correctly applies a 2D inverse rotation
+    matrix to convert the global target vector into a local (egocentric) vector.
+    """
+    cfg = BrittleStarConfig()
+    env = BrittleStarJaxEnvWrapper.default(num_envs=1, backend=Backend.MJX)
+
+    obs_processor = create_obs_processor(
+        bounds_dict=cfg.obs_bounds.to_bounds_dict(), padding_masks=env.padding_masks
+    )
+
+    env_state = env.reset(seed=42)
+
+    # --- Scenario 1 ---
+    # Robot is rotated 90 degrees Left (facing global Y)
+    # Target is straight ahead on the global X axis [1.0, 0.0]
+    # Because the robot is facing Y, the target on X is to its RIGHT [0.0, -1.0] locally.
+    dummy_obs_1 = dict(env_state.observations)
+    dummy_obs_1["disk_rotation"] = jnp.array([[0.0, 0.0, jnp.pi / 2.0]])
+    dummy_obs_1["unit_xy_direction_to_target"] = jnp.array([[1.0, 0.0]])
+    processed_1 = obs_processor(dummy_obs_1)
+
+    # --- Scenario 2 (used to find the array indices) ---
+    # We change ONLY the target vector so we can isolate it in the final array
+    dummy_obs_2 = dict(env_state.observations)
+    dummy_obs_2["disk_rotation"] = jnp.array([[0.0, 0.0, jnp.pi / 2.0]])
+    dummy_obs_2["unit_xy_direction_to_target"] = jnp.array([[0.0, 1.0]])
+    processed_2 = obs_processor(dummy_obs_2)
+
+    # Find the indices of the elements that changed
+    diff_array = jnp.abs(processed_1[0] - processed_2[0])
+    changed_indices = jnp.where(diff_array > 1e-4)[0]
+
+    local_target = processed_1[0, changed_indices]
+    expected_local_target = jnp.array([0.0, -1.0])
+
+    assert jnp.sum(jnp.abs(local_target - expected_local_target)) < 1e-4, (
+        f"The obs_processor did not correctly rotate the vector to egocentric. "
+        f"Expected {expected_local_target}, but got {local_target}."
+    )