<gh_stars>0
import numpy as np


class Reward:
    """A simple wrapper class for reward computation."""
    def __call__(self, trajectories, shapes):
        raise NotImplementedError

    @property
    def name(self):
        return self.NAME


class ZeroReward(Reward):
    """Returns zero reward for every shape and step in trajectory."""
    NAME = 'zero'

    def __call__(self, trajectories, shapes):
        return np.zeros((trajectories.shape[0],), dtype=np.float32)


class VertPosReward(Reward):
    """Returns reward proportional to the vertical position of the first object."""
    NAME = 'vertical_position'

    def __call__(self, trajectories, shapes):
        return trajectories[:, 0, 1]


class HorPosReward(Reward):
    """Returns reward proportional to the horizontal position of the first object."""
    NAME = 'horizontal_position'

    def __call__(self, trajectories, shapes):
        return trajectories[:, 0, 0]


class AgentXReward(Reward):
    """Returns reward proportional to the horizontal position of the agent. Assumes that agent is the first object."""
    NAME = 'agent_x'

    def __call__(self, trajectories, shapes):
        return trajectories[:, 0, 1]


class AgentYReward(Reward):
    """Returns reward proportional to the vertical position of the agent. Assumes that agent is the first object."""
    NAME = 'agent_y'

    def __call__(self, trajectories, shapes):
        return trajectories[:, 0, 0]


class TargetXReward(Reward):
    """Returns reward proportional to the horizontal position of the target. Assumes that target is second object."""
    NAME = 'target_x'

    def __call__(self, trajectories, shapes):
        return trajectories[:, 1, 1]


class TargetYReward(Reward):
    """Returns reward proportional to the vertical position of the target. Assumes that target is the second object."""
    NAME = 'target_y'

    def __call__(self, trajectories, shapes):
        return trajectories[:, 1, 0]


class FollowReward(Reward):
    """Returns reward (inversely) proportional to the distance between the agent and the target."""
    NAME = 'follow'

    def __call__(self, trajectories, shapes):
        agent_pos = trajectories[:, 0, :2]
        target_pos = trajectories[:, 1, :2]

        reward = [1. - np.sqrt(((target_pos[i] - agent_pos[i]) ** 2).sum()) /
                  np.sqrt(2) for i in range(trajectories.shape[0])]
        return np.asarray(reward, dtype=np.float32)

import numpy as np


class Reward:
    """A simple wrapper class for reward computation."""
    def __call__(self, trajectories, shapes):
        raise NotImplementedError

    @property
    def name(self):
        return self.NAME


class ZeroReward(Reward):
    """Returns zero reward for every shape and step in trajectory."""
    NAME = 'zero'

    def __call__(self, trajectories, shapes):
        return np.zeros((trajectories.shape[0],), dtype=np.float32)


class VertPosReward(Reward):
    """Returns reward proportional to the vertical position of the first object."""
    NAME = 'vertical_position'

    def __call__(self, trajectories, shapes):
        return trajectories[:, 0, 1]


class HorPosReward(Reward):
    """Returns reward proportional to the horizontal position of the first object."""
    NAME = 'horizontal_position'

    def __call__(self, trajectories, shapes):
        return trajectories[:, 0, 0]


class AgentXReward(Reward):
    """Returns reward proportional to the horizontal position of the agent. Assumes that agent is the first object."""
    NAME = 'agent_x'

    def __call__(self, trajectories, shapes):
        return trajectories[:, 0, 1]


class AgentYReward(Reward):
    """Returns reward proportional to the vertical position of the agent. Assumes that agent is the first object."""
    NAME = 'agent_y'

    def __call__(self, trajectories, shapes):
        return trajectories[:, 0, 0]


class TargetXReward(Reward):
    """Returns reward proportional to the horizontal position of the target. Assumes that target is second object."""
    NAME = 'target_x'

    def __call__(self, trajectories, shapes):
        return trajectories[:, 1, 1]


class TargetYReward(Reward):
    """Returns reward proportional to the vertical position of the target. Assumes that target is the second object."""
    NAME = 'target_y'

    def __call__(self, trajectories, shapes):
        return trajectories[:, 1, 0]


class FollowReward(Reward):
    """Returns reward (inversely) proportional to the distance between the agent and the target."""
    NAME = 'follow'

    def __call__(self, trajectories, shapes):
        agent_pos = trajectories[:, 0, :2]
        target_pos = trajectories[:, 1, :2]

        reward = [1. - np.sqrt(((target_pos[i] - agent_pos[i]) ** 2).sum()) /
                  np.sqrt(2) for i in range(trajectories.shape[0])]
        return np.asarray(reward, dtype=np.float32)
A simple wrapper class for reward computation. Returns zero reward for every shape and step in trajectory. Returns reward proportional to the vertical position of the first object. Returns reward proportional to the horizontal position of the first object. Returns reward proportional to the horizontal position of the agent. Assumes that agent is the first object. Returns reward proportional to the vertical position of the agent. Assumes that agent is the first object. Returns reward proportional to the horizontal position of the target. Assumes that target is second object. Returns reward proportional to the vertical position of the target. Assumes that target is the second object. Returns reward (inversely) proportional to the distance between the agent and the target.