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