fix: custom reward function dependent on env reward + extensions
This commit is contained in:
Robin Meersman
parent
3efee5d746
commit
ad17401c4c
9 changed files with 71 additions and 50 deletions
|
|
@ -27,7 +27,7 @@ class BrittleStarJaxEnvWrapper:
|
|||
)
|
||||
|
||||
# Pre-compute masks for observation padding
|
||||
self._padding_masks = compute_padding_masks(self._morphology.segments_per_arm)
|
||||
self._padding_masks = compute_padding_masks(self._morphology.segments_per_arm, (4, 4))
|
||||
|
||||
self._vectorized_reset = jax.jit(jax.vmap(self._env.reset))
|
||||
self._vectorized_step = jax.jit(jax.vmap(self._env.step))
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
@ -630,8 +632,6 @@ class PPOTrainer:
|
|||
iteration_time_start,
|
||||
training_measurements,
|
||||
storage,
|
||||
next_obs,
|
||||
xy_distance,
|
||||
)
|
||||
|
||||
sps = int(global_step / (time.time() - start_time))
|
||||
|
|
|
|||
Reference in a new issue