| import numpy as np | |
| def shape_reward_vectorized(states, raw_rewards): | |
| """ | |
| Shapes rewards for an array of states originating from multiple environments. | |
| states shape: (NUM_ENVS, 2) | |
| raw_rewards shape: (NUM_ENVS,) | |
| """ | |
| positions = states[:, 0] | |
| velocities = states[:, 1] | |
| velocity_bonus = np.abs(velocities) * 10.0 | |
| height_bonus = (positions + 0.5) ** 2 | |
| return raw_rewards + velocity_bonus + height_bonus |