tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
Archived
1
Fork 0
Code Issues Releases Activity

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

Browse source
This commit is contained in:
RobinMeersman 2026-04-18 15:52:33 +02:00 • committed by GitHub
commit ea59c821a3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 71 additions and 69 deletions
Download patch file Download diff file Expand all files Collapse all files

4
configs/environment/directed_locomotion.yaml
View file

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

5
configs/morphology/2_arms.yaml Normal file
View file

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

2
configs/ppo/debug.yaml
View file

@ -1,5 +1,5 @@
learning_rate: 0.0003
total_timesteps: 409600
total_timesteps: 409600
num_envs: 32
num_steps: 32
anneal_lr: true

16
configs/ppo/dev_larger_timesteps_larger_rolloutsteps.yaml Normal file
View file

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

19
configs/ppo/fast.yaml
View file

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

2
docs/design/reward_function.md
View file

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

2
env/hpc/requirements.txt vendored
View file

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

6
src/brittle_star_project/environment/padded_obs_wrapper.py
View file

@ -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, ...],
reference_segments_per_arm: tuple[int, ...] = (4, 4, 4, 4, 4),
segments_per_arm: Sequence[int],
reference_segments_per_arm: Sequence[int] = (4, 4, 4, 4, 4),
) -> dict[str, Any]:
"""Pre-compute boolean masks for spatial insertion of observations.

84
src/brittle_star_project/trainers/PPOTrainer.py
View file

@ -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 *= 20000
reward = jnp.clip(reward, -10, 10)
reward = _reward_fn(env_state, next_env_state)
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],
"values": storage.values[0],
"returns": storage.returns[0],
"advantages": storage.advantages[0],
"actions": storage.actions[0],
"raw_actions": storage.raw_actions[0],
"means": storage.means[0],
"stds": storage.stds[0],
"logprobs": storage.logprobs[0],
"rewards": storage.rewards,
"values": storage.values,
"returns": storage.returns,
"advantages": storage.advantages,
}
)
storage_metrics = {
"rollout/env0/return_mean": float(np.mean(data["returns"])),
"rollout/env0/advantage_mean": float(np.mean(data["advantages"])),
"rollout/env0/value_mean": float(np.mean(data["values"])),
"rollout/env0/value_vs_return_diff": float(np.mean(data["values"] - data["returns"])),
"rollout/env0/reward_mean": float(np.mean(data["rewards"])),
"rollout/env0/mean_mean": float(np.mean(data["means"])),
"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"])),
rollout_metrics = {
"rollout/reward_mean": float(np.mean(data["rewards"])),
"rollout/return_mean": float(np.mean(data["returns"])),
"rollout/value_mean": float(np.mean(data["values"])),
"rollout/advantage_mean": float(np.mean(data["advantages"])),
"rollout/advantage_std": float(np.std(data["advantages"])),
"rollout/value_vs_return_mse": float(np.mean((data["values"] - data["returns"]) ** 2)),
}
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/avg_episodic_length": np.mean(
jax.device_get(episode_stats.returned_episode_lengths)
"charts/episodic_return": training_measurements.avg_episodic_return,
"charts/episodic_length": float(
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))