feat: improved the used reward function to better the training results

configs/environment/directed_locomotion.yaml

@@ -2,11 +2,11 @@
 # Baseline task setting.
 
 task: DIRECTED_LOCOMOTION
-simulation_time: 5.0
+simulation_time: 5000.0
 num_physics_steps_per_control_step: 10
 time_scale: 2
 camera_ids: [0, 1]
 render_size: [480, 640]
 joint_randomization_noise_scale: 0.0
-target_distance: 3.0
+target_distance: 0.6
 light_perlin_noise_scale: 0

configs/morphology/2_arms.yaml

@@ -0,0 +1,5 @@
+# 2 Arms Morphology Configuration
+
+segments_per_arm: [4, 0, 4, 0, 0]
+use_p_control: true
+use_torque_control: false

configs/ppo/dev_larger_timesteps_larger_rolloutsteps.yaml

@@ -0,0 +1,16 @@
+learning_rate: 0.0003
+total_timesteps: 1228800
+num_envs: 32
+num_steps: 64
+anneal_lr: true
+gamma: 0.99
+gae_lambda: 0.95
+num_minibatches: 32
+update_epochs: 4
+norm_adv: true
+clip_coef: 0.2
+clip_vloss: true
+ent_coef: 0.005
+vf_coef: 1.0
+max_grad_norm: 0.5
+target_kl: null

configs/ppo/fast.yaml

@@ -1,19 +0,0 @@
-# Fast PPO Configuration
-# Lower timestep count for quick iterations/testing.
-
-learning_rate: 0.0005
-total_timesteps: 500000
-num_envs: 8
-num_steps: 128
-anneal_lr: true
-gamma: 0.99
-gae_lambda: 0.95
-num_minibatches: 4
-update_epochs: 4
-norm_adv: true
-clip_coef: 0.2
-clip_vloss: true
-ent_coef: 0.01
-vf_coef: 0.5
-max_grad_norm: 0.5
-target_kl: null

docs/design/reward_function.md

@@ -8,6 +8,8 @@ inputs must be distributed fairly to guarantee an objective comparison between d
 - The reward function is centered around minimizing the distance to the goal or maximizing the movement towards the goal
   within a finite number of timesteps $T$.
 - To motivate efficient movement, the amount of timesteps taken to reach the goal will be used as penalty.
+- An extra penalty based on movement relative to the current step and
+the previous is used to penalize a movement away from the target.
 
 ## From reward to PPO
 

env/hpc/requirements.txt

@@ -15,6 +15,6 @@ optax>=0.2.6
 pyopengl>=3.1.10
 pyopengl-accelerate>=3.1.10
 pyyaml>=6.0
-tyro>=1.0.10
+hydra-core>=1.3.2
 wandb==0.24.2
 torch>=2.4.0

src/brittle_star_project/environment/padded_obs_wrapper.py

@@ -8,7 +8,7 @@ flattened observation maintains the correct physical mapping to the neural netwo
 
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, Sequence
 import jax.numpy as jnp
 
 # Observation keys whose size scales with the number of joints (2 per segment).
@@ -30,8 +30,8 @@ _SEGMENT_SCALED_KEYS = frozenset(
 
 
 def compute_padding_masks(
-    segments_per_arm: tuple[int, ...],
+    segments_per_arm: Sequence[int],
-    reference_segments_per_arm: tuple[int, ...] = (4, 4, 4, 4, 4),
+    reference_segments_per_arm: Sequence[int] = (4, 4, 4, 4, 4),
 ) -> dict[str, Any]:
     """Pre-compute boolean masks for spatial insertion of observations.
 

src/brittle_star_project/trainers/PPOTrainer.py

@@ -151,12 +151,27 @@ def _step_once(
     return (agent_state, episode_stats, next_obs, next_done, key, env_state), storage
 
 
+def _reward_fn(env_state, next_env_state):
+    # if delta distance positive ==> brittle star walking away from target
+    delta_distance = (
+        next_env_state.observations["xy_distance_to_target"]
+        - env_state.observations["xy_distance_to_target"]
+    ).squeeze(-1)
+
+    env_reward = next_env_state.reward
+    clipped_env_reward = jnp.clip(100 * env_reward, -10, 10)
+
+    time_penalty = 0.1
+    distance_penalty = jnp.clip(0.5 * delta_distance, -0.5, 0.5)
+    penalty = time_penalty + distance_penalty
+
+    return jnp.where(next_env_state.terminated, 50.0, clipped_env_reward - penalty)
+
+
 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 = _reward_fn(env_state, next_env_state)
-    reward *= 20000
-    reward = jnp.clip(reward, -10, 10)
     terminated = next_env_state.terminated
     truncated = next_env_state.truncated
     done = terminated | truncated
@@ -440,62 +455,49 @@ class PPOTrainer:
         iteration_time_start,
         training_measurements,
         storage,
-        next_obs,
-        xy_distance,
     ):
         data = jax.device_get(
             {
-                "rewards": storage.rewards[0],
+                "rewards": storage.rewards,
-                "values": storage.values[0],
+                "values": storage.values,
-                "returns": storage.returns[0],
+                "returns": storage.returns,
-                "advantages": storage.advantages[0],
+                "advantages": storage.advantages,
-                "actions": storage.actions[0],
-                "raw_actions": storage.raw_actions[0],
-                "means": storage.means[0],
-                "stds": storage.stds[0],
-                "logprobs": storage.logprobs[0],
             }
         )
 
-        storage_metrics = {
+        rollout_metrics = {
-            "rollout/env0/return_mean": float(np.mean(data["returns"])),
+            "rollout/reward_mean": float(np.mean(data["rewards"])),
-            "rollout/env0/advantage_mean": float(np.mean(data["advantages"])),
+            "rollout/return_mean": float(np.mean(data["returns"])),
-            "rollout/env0/value_mean": float(np.mean(data["values"])),
+            "rollout/value_mean": float(np.mean(data["values"])),
-            "rollout/env0/value_vs_return_diff": float(np.mean(data["values"] - data["returns"])),
+            "rollout/advantage_mean": float(np.mean(data["advantages"])),
-            "rollout/env0/reward_mean": float(np.mean(data["rewards"])),
+            "rollout/advantage_std": float(np.std(data["advantages"])),
-            "rollout/env0/mean_mean": float(np.mean(data["means"])),
+            "rollout/value_vs_return_mse": float(np.mean((data["values"] - data["returns"]) ** 2)),
-            "rollout/env0/logprob_mean": float(np.mean(data["logprobs"])),
-            "rollout/env0/action_mean": float(np.mean(data["actions"])),
-            "rollout/env0/raw_action_mean": float(np.mean(data["raw_actions"])),
         }
 
-        for i in range(len(xy_distance)):
-            storage_metrics[f"env_data/env{i}_xy_dist_target"] = float(xy_distance[i])
-
         metrics = {
-            "charts/avg_episodic_return": training_measurements.avg_episodic_return,
+            "charts/episodic_return": training_measurements.avg_episodic_return,
-            "charts/avg_episodic_length": np.mean(
+            "charts/episodic_length": float(
-                jax.device_get(episode_stats.returned_episode_lengths)
+                np.mean(jax.device_get(episode_stats.returned_episode_lengths))
             ),
-            "charts/learning_rate": self.agent_state.opt_state[1]
-            .hyperparams["learning_rate"]
-            .item(),
             "charts/explained_variance": training_measurements.explained_variance,
-            "charts/num_terminated": training_measurements.num_terminated,
-            "charts/num_truncated": training_measurements.num_truncated,
-            "charts/avg_terminated_ep_length": training_measurements.avg_terminated_length,
-            "charts/avg_truncated_ep_length": training_measurements.avg_truncated_length,
             "losses/value_loss": training_measurements.v_loss[-1, -1].item(),
             "losses/policy_loss": training_measurements.pg_loss[-1, -1].item(),
             "losses/entropy": training_measurements.entropy_loss[-1, -1].item(),
             "losses/approx_kl": training_measurements.approx_kl[-1, -1].item(),
-            "losses/loss": training_measurements.loss[-1, -1].item(),
+            "charts/learning_rate": self.agent_state.opt_state[1]
+            .hyperparams["learning_rate"]
+            .item(),
             "charts/SPS": int(global_step / (time.time() - start_time)),
             "charts/SPS_update": int(
                 self.ppo.num_envs * self.ppo.num_steps / (time.time() - iteration_time_start)
             ),
-            **storage_metrics,
+            "termi_trunci/num_terminated": training_measurements.num_terminated,
+            "termi_trunci/num_truncated": training_measurements.num_truncated,
+            "termi_trunci/avg_terminated_ep_length": training_measurements.avg_terminated_length,
+            "termi_trunci/avg_truncated_ep_length": training_measurements.avg_truncated_length,
+            **rollout_metrics,
         }
+
         self.logger.log(metrics, step=global_step)
 
     def _step(self, env_state, next_obs, next_done, iteration: int) -> tuple:
@@ -610,8 +612,6 @@ class PPOTrainer:
             self._update_obs_stats(next_obs)
             next_obs = _normalize_obs(next_obs, self.obs_mean, self.obs_var)
 
-            xy_distance = _get_xy_distance_to_target(env_state.observations)
-
             global_step += self.ppo.num_steps * self.ppo.num_envs
             self._log(
                 global_step,
@@ -620,8 +620,6 @@ class PPOTrainer:
                 iteration_time_start,
                 training_measurements,
                 storage,
-                next_obs,
-                xy_distance,
             )
 
             sps = int(global_step / (time.time() - start_time))