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

feat: modded reward?

Browse source
This commit is contained in:
JibrilExe 2026-04-16 09:41:10 +02:00
parent 467483d0ae
commit b48aea619a
1 changed files with 13 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -27,10 +27,10 @@ from experiment_logger import get_logger
_ALLOWED_OBS_KEYS = { _ALLOWED_OBS_KEYS = {
"joint_position", "joint_position",
"joint_velocity", "joint_velocity",
"joint_actuator_force", # "joint_actuator_force",
"actuator_force", # "actuator_force",
"disk_position", "disk_position",
"disk_rotation", # "disk_rotation",
"disk_linear_velocity", "disk_linear_velocity",
"disk_angular_velocity", "disk_angular_velocity",
"unit_xy_direction_to_target", "unit_xy_direction_to_target",
@ -153,7 +153,10 @@ def _step_env_wrapped(episode_stats, env_state, action, env_step_fn):
next_env_state = env_step_fn(env_state, action) next_env_state = env_step_fn(env_state, action)
reward = next_env_state.reward reward = next_env_state.reward
reward *= 20000 reward *= 100
xy_dist = _get_xy_distance_to_target(next_env_state.observations) # shape: (num_envs,)
distance_scale = 1
reward = reward - distance_scale * xy_dist
reward = jnp.clip(reward, -10, 10) reward = jnp.clip(reward, -10, 10)
terminated = next_env_state.terminated terminated = next_env_state.terminated
truncated = next_env_state.truncated truncated = next_env_state.truncated
@ -434,7 +437,7 @@ class PPOTrainer:
): ):
data = jax.device_get( data = jax.device_get(
{ {
"rewards": storage.rewards[0], "rewards": storage.rewards,
"values": storage.values[0], "values": storage.values[0],
"returns": storage.returns[0], "returns": storage.returns[0],
"advantages": storage.advantages[0], "advantages": storage.advantages[0],
@ -458,8 +461,12 @@ class PPOTrainer:
"rollout/env0/raw_action_mean": float(np.mean(data["raw_actions"])), "rollout/env0/raw_action_mean": float(np.mean(data["raw_actions"])),
} }
rewards_per_env = np.mean(data["rewards"], axis=0) # mean reward per env over the rollout steps
for i in range(len(xy_distance)): for i in range(len(xy_distance)):
storage_metrics[f"env_data/env{i}_xy_dist_target"] = float(xy_distance[i]) env_prefix = f"env_data{i}"
storage_metrics[f"{env_prefix}/reward_mean"] = float(rewards_per_env[i])
storage_metrics[f"{env_prefix}/xy_dist_target"] = float(xy_distance[i])
storage_metrics[f"{env_prefix}/last_reward"] = float(data["rewards"][-1, i])
metrics = { metrics = {
"charts/avg_episodic_return": training_measurements.avg_episodic_return, "charts/avg_episodic_return": training_measurements.avg_episodic_return,