<reponame>FineArtz/ILSwiss<gh_stars>10-100
import numpy as np
from rlkit.scripted_experts.scripted_policy import ScriptedPolicy


class ScriptedMultiDirectionPointMassPolicy(ScriptedPolicy):
    def __init__(self):
        self._ptr = 0
        self.angles = np.array(
            [
                0,
                np.pi / 4.0,
                np.pi / 2.0,
                3 * np.pi / 4.0,
                np.pi,
                5 * np.pi / 4.0,
                3 * np.pi / 2.0,
                7 * np.pi / 4.0,
            ]
        )

    def reset(self, env):
        # a = np.random.uniform(np.pi/6.0, 5*np.pi/6.0)
        # a = self._ptr*(4*np.pi/6.0)/300 + np.pi/6.0
        a = self.angles[self._ptr]
        self._ptr += 1
        self._ptr %= self.angles.shape[0]
        # self.target = 100.0*np.array([np.cos(a), np.sin(a)])
        self.target = 25.0 * np.array([np.cos(a), np.sin(a)])

    def get_action(self, obs, env, timestep):
        act = self.target - env.cur_pos
        act = np.random.uniform(0.75, 0.95) * act / np.linalg.norm(act)
        return act, {}
