from space import Box
import numpy as np
import seeding

class Env(object):
    """The main OpenAI Gym class. It encapsulates an environment with
    arbitrary behind-the-scenes dynamics. An environment can be
    partially or fully observed.

    The main API methods that users of this class need to know are:

        step
        reset
        render
        close
        seed

    And set the following attributes:

        action_space: The Space object corresponding to valid actions
        observation_space: The Space object corresponding to valid observations
        reward_range: A tuple corresponding to the min and max possible rewards

    Note: a default reward range set to [-inf,+inf] already exists. Set it if you want a narrower range.

    The methods are accessed publicly as "step", "reset", etc...
    """
    # Set this in SOME subclasses
    metadata = {'render.modes': []}
    reward_range = (-float('inf'), float('inf'))
    spec = None

    # Set these in ALL subclasses
    action_space = None
    observation_space = None

    def step(self, action):
        """Run one timestep of the environment's dynamics. When end of
        episode is reached, you are responsible for calling `reset()`
        to reset this environment's state.

        Accepts an action and returns a tuple (observation, reward, done, info).

        Args:
            action (object): an action provided by the agent

        Returns:
            observation (object): agent's observation of the current environment
            reward (float) : amount of reward returned after previous action
            done (bool): whether the episode has ended, in which case further step() calls will return undefined results
            info (dict): contains auxiliary diagnostic information (helpful for debugging, and sometimes learning)
        """
        raise NotImplementedError

    def reset(self):
        """Resets the environment to an initial state and returns an initial
        observation.

        Note that this function should not reset the environment's random
        number generator(s); random variables in the environment's state should
        be sampled independently between multiple calls to `reset()`. In other
        words, each call of `reset()` should yield an environment suitable for
        a new episode, independent of previous episodes.

        Returns:
            observation (object): the initial observation.
        """
        raise NotImplementedError

    def render(self, mode='human'):
        """Renders the environment.

        The set of supported modes varies per environment. (And some
        environments do not support rendering at all.) By convention,
        if mode is:

        - human: render to the current display or terminal and
          return nothing. Usually for human consumption.
        - rgb_array: Return an numpy.ndarray with shape (x, y, 3),
          representing RGB values for an x-by-y pixel image, suitable
          for turning into a video.
        - ansi: Return a string (str) or StringIO.StringIO containing a
          terminal-style text representation. The text can include newlines
          and ANSI escape sequences (e.g. for colors).

        Note:
            Make sure that your class's metadata 'render.modes' key includes
              the list of supported modes. It's recommended to call super()
              in implementations to use the functionality of this method.

        Args:
            mode (str): the mode to render with

        Example:

        class MyEnv(Env):
            metadata = {'render.modes': ['human', 'rgb_array']}

            def render(self, mode='human'):
                if mode == 'rgb_array':
                    return np.array(...) # return RGB frame suitable for video
                elif mode == 'human':
                    ... # pop up a window and render
                else:
                    super(MyEnv, self).render(mode=mode) # just raise an exception
        """
        raise NotImplementedError

    def close(self):
        """Override close in your subclass to perform any necessary cleanup.

        Environments will automatically close() themselves when
        garbage collected or when the program exits.
        """
        pass

    def seed(self, seed=None):
        """Sets the seed for this env's random number generator(s).

        Note:
            Some environments use multiple pseudorandom number generators.
            We want to capture all such seeds used in order to ensure that
            there aren't accidental correlations between multiple generators.

        Returns:
            list<bigint>: Returns the list of seeds used in this env's random
              number generators. The first value in the list should be the
              "main" seed, or the value which a reproducer should pass to
              'seed'. Often, the main seed equals the provided 'seed', but
              this won't be true if seed=None, for example.
        """
        return

    @property
    def unwrapped(self):
        """Completely unwrap this env.

        Returns:
            gym.Env: The base non-wrapped gym.Env instance
        """
        return self

    def __str__(self):
        if self.spec is None:
            return '<{} instance>'.format(type(self).__name__)
        else:
            return '<{}<{}>>'.format(type(self).__name__, self.spec.id)

    def __enter__(self):
        """Support with-statement for the environment. """
        return self

    def __exit__(self, *args):
        """Support with-statement for the environment. """
        self.close()
        # propagate exception
        return False

class PendulumEnv(Env):
    metadata = {
        'render.modes': ['human', 'rgb_array'],
        'video.frames_per_second': 30
    }

    def __init__(self, g=9.8):
        self.max_speed = 2 * np.pi
        self.max_torque = 5.0
        self.dt = .1
        self.g = g
        self.m = 1.
        self.l = 1.
        self.mu = 0.1
        self.viewer = None

        high = np.array([np.pi, self.max_speed], dtype=np.float32)
        self.action_space = Box(
            low=-self.max_torque,
            high=self.max_torque, shape=(1,),
            dtype=np.float32
        )
        self.observation_space = Box(
            low=-high,
            high=high,
            dtype=np.float32
        )

        self.seed()

    def seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]

    def step(self, u):
        th, thdot = self.state  # th := theta

        g = self.g
        m = self.m
        l = self.l
        mu = self.mu
        dt = self.dt / 100 # since the euler method used for dt = 0.001

        u = np.clip(u, -self.max_torque, self.max_torque)[0]
        self.last_u = u  # for rendering
        costs = np.abs(angle_normalize(th))
        for i in range(100):
            newthdot = (1 - dt * mu / (m * l**2)) * thdot + (g / l * np.sin(th) + u / (m * l ** 2)) * dt
            newth = th + newthdot * dt
            newthdot = np.clip(newthdot, -self.max_speed, self.max_speed)
            th = newth
            thdot = newthdot
        newth = angle_normalize(newth)
        self.state = np.array([newth, newthdot])
        return self.state, -costs, False, {}

    def reset(self):
        self.state = [np.pi, 2*np.pi]
        self.last_u = None
        return self.state

    def render(self, mode='human'):
        if self.viewer is None:
            import rendering
            self.viewer = rendering.Viewer(500, 500)
            self.viewer.set_bounds(-2.2, 2.2, -2.2, 2.2)
            rod = rendering.make_capsule(1, .2)
            rod.set_color(.8, .3, .3)
            self.pole_transform = rendering.Transform()
            rod.add_attr(self.pole_transform)
            self.viewer.add_geom(rod)
            axle = rendering.make_circle(.05)
            axle.set_color(0, 0, 0)
            self.viewer.add_geom(axle)
        self.pole_transform.set_rotation(self.state[0])

        return self.viewer.render(return_rgb_array=mode == 'rgb_array')

    def close(self):
        if self.viewer:
            self.viewer.close()
            self.viewer = None


def angle_normalize(x):
    return (((x+np.pi) % (2*np.pi)) - np.pi)
from space import Box
import numpy as np
import seeding

class Env(object):
    """The main OpenAI Gym class. It encapsulates an environment with
    arbitrary behind-the-scenes dynamics. An environment can be
    partially or fully observed.

    The main API methods that users of this class need to know are:

        step
        reset
        render
        close
        seed

    And set the following attributes:

        action_space: The Space object corresponding to valid actions
        observation_space: The Space object corresponding to valid observations
        reward_range: A tuple corresponding to the min and max possible rewards

    Note: a default reward range set to [-inf,+inf] already exists. Set it if you want a narrower range.

    The methods are accessed publicly as "step", "reset", etc...
    """
    # Set this in SOME subclasses
    metadata = {'render.modes': []}
    reward_range = (-float('inf'), float('inf'))
    spec = None

    # Set these in ALL subclasses
    action_space = None
    observation_space = None

    def step(self, action):
        """Run one timestep of the environment's dynamics. When end of
        episode is reached, you are responsible for calling `reset()`
        to reset this environment's state.

        Accepts an action and returns a tuple (observation, reward, done, info).

        Args:
            action (object): an action provided by the agent

        Returns:
            observation (object): agent's observation of the current environment
            reward (float) : amount of reward returned after previous action
            done (bool): whether the episode has ended, in which case further step() calls will return undefined results
            info (dict): contains auxiliary diagnostic information (helpful for debugging, and sometimes learning)
        """
        raise NotImplementedError

    def reset(self):
        """Resets the environment to an initial state and returns an initial
        observation.

        Note that this function should not reset the environment's random
        number generator(s); random variables in the environment's state should
        be sampled independently between multiple calls to `reset()`. In other
        words, each call of `reset()` should yield an environment suitable for
        a new episode, independent of previous episodes.

        Returns:
            observation (object): the initial observation.
        """
        raise NotImplementedError

    def render(self, mode='human'):
        """Renders the environment.

        The set of supported modes varies per environment. (And some
        environments do not support rendering at all.) By convention,
        if mode is:

        - human: render to the current display or terminal and
          return nothing. Usually for human consumption.
        - rgb_array: Return an numpy.ndarray with shape (x, y, 3),
          representing RGB values for an x-by-y pixel image, suitable
          for turning into a video.
        - ansi: Return a string (str) or StringIO.StringIO containing a
          terminal-style text representation. The text can include newlines
          and ANSI escape sequences (e.g. for colors).

        Note:
            Make sure that your class's metadata 'render.modes' key includes
              the list of supported modes. It's recommended to call super()
              in implementations to use the functionality of this method.

        Args:
            mode (str): the mode to render with

        Example:

        class MyEnv(Env):
            metadata = {'render.modes': ['human', 'rgb_array']}

            def render(self, mode='human'):
                if mode == 'rgb_array':
                    return np.array(...) # return RGB frame suitable for video
                elif mode == 'human':
                    ... # pop up a window and render
                else:
                    super(MyEnv, self).render(mode=mode) # just raise an exception
        """
        raise NotImplementedError

    def close(self):
        """Override close in your subclass to perform any necessary cleanup.

        Environments will automatically close() themselves when
        garbage collected or when the program exits.
        """
        pass

    def seed(self, seed=None):
        """Sets the seed for this env's random number generator(s).

        Note:
            Some environments use multiple pseudorandom number generators.
            We want to capture all such seeds used in order to ensure that
            there aren't accidental correlations between multiple generators.

        Returns:
            list<bigint>: Returns the list of seeds used in this env's random
              number generators. The first value in the list should be the
              "main" seed, or the value which a reproducer should pass to
              'seed'. Often, the main seed equals the provided 'seed', but
              this won't be true if seed=None, for example.
        """
        return

    @property
    def unwrapped(self):
        """Completely unwrap this env.

        Returns:
            gym.Env: The base non-wrapped gym.Env instance
        """
        return self

    def __str__(self):
        if self.spec is None:
            return '<{} instance>'.format(type(self).__name__)
        else:
            return '<{}<{}>>'.format(type(self).__name__, self.spec.id)

    def __enter__(self):
        """Support with-statement for the environment. """
        return self

    def __exit__(self, *args):
        """Support with-statement for the environment. """
        self.close()
        # propagate exception
        return False

class PendulumEnv(Env):
    metadata = {
        'render.modes': ['human', 'rgb_array'],
        'video.frames_per_second': 30
    }

    def __init__(self, g=9.8):
        self.max_speed = 2 * np.pi
        self.max_torque = 5.0
        self.dt = .1
        self.g = g
        self.m = 1.
        self.l = 1.
        self.mu = 0.1
        self.viewer = None

        high = np.array([np.pi, self.max_speed], dtype=np.float32)
        self.action_space = Box(
            low=-self.max_torque,
            high=self.max_torque, shape=(1,),
            dtype=np.float32
        )
        self.observation_space = Box(
            low=-high,
            high=high,
            dtype=np.float32
        )

        self.seed()

    def seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]

    def step(self, u):
        th, thdot = self.state  # th := theta

        g = self.g
        m = self.m
        l = self.l
        mu = self.mu
        dt = self.dt / 100 # since the euler method used for dt = 0.001

        u = np.clip(u, -self.max_torque, self.max_torque)[0]
        self.last_u = u  # for rendering
        costs = np.abs(angle_normalize(th))
        for i in range(100):
            newthdot = (1 - dt * mu / (m * l**2)) * thdot + (g / l * np.sin(th) + u / (m * l ** 2)) * dt
            newth = th + newthdot * dt
            newthdot = np.clip(newthdot, -self.max_speed, self.max_speed)
            th = newth
            thdot = newthdot
        newth = angle_normalize(newth)
        self.state = np.array([newth, newthdot])
        return self.state, -costs, False, {}

    def reset(self):
        self.state = [np.pi, 2*np.pi]
        self.last_u = None
        return self.state

    def render(self, mode='human'):
        if self.viewer is None:
            import rendering
            self.viewer = rendering.Viewer(500, 500)
            self.viewer.set_bounds(-2.2, 2.2, -2.2, 2.2)
            rod = rendering.make_capsule(1, .2)
            rod.set_color(.8, .3, .3)
            self.pole_transform = rendering.Transform()
            rod.add_attr(self.pole_transform)
            self.viewer.add_geom(rod)
            axle = rendering.make_circle(.05)
            axle.set_color(0, 0, 0)
            self.viewer.add_geom(axle)
        self.pole_transform.set_rotation(self.state[0])

        return self.viewer.render(return_rgb_array=mode == 'rgb_array')

    def close(self):
        if self.viewer:
            self.viewer.close()
            self.viewer = None


def angle_normalize(x):
    return (((x+np.pi) % (2*np.pi)) - np.pi)
The main OpenAI Gym class. It encapsulates an environment with     arbitrary behind-the-scenes dynamics. An environment can be     partially or fully observed.      The main API methods that users of this class need to know are:          step         reset         render         close         seed      And set the following attributes:          action_space: The Space object corresponding to valid actions         observation_space: The Space object corresponding to valid observations         reward_range: A tuple corresponding to the min and max possible rewards      Note: a default reward range set to [-inf,+inf] already exists. Set it if you want a narrower range.      The methods are accessed publicly as "step", "reset", etc... # Set this in SOME subclasses # Set these in ALL subclasses Run one timestep of the environment's dynamics. When end of         episode is reached, you are responsible for calling `reset()`         to reset this environment's state.          Accepts an action and returns a tuple (observation, reward, done, info).          Args:             action (object): an action provided by the agent          Returns:             observation (object): agent's observation of the current environment             reward (float) : amount of reward returned after previous action             done (bool): whether the episode has ended, in which case further step() calls will return undefined results             info (dict): contains auxiliary diagnostic information (helpful for debugging, and sometimes learning) Resets the environment to an initial state and returns an initial         observation.          Note that this function should not reset the environment's random         number generator(s); random variables in the environment's state should         be sampled independently between multiple calls to `reset()`. In other         words, each call of `reset()` should yield an environment suitable for         a new episode, independent of previous episodes.          Returns:             observation (object): the initial observation. Renders the environment.          The set of supported modes varies per environment. (And some         environments do not support rendering at all.) By convention,         if mode is:          - human: render to the current display or terminal and           return nothing. Usually for human consumption.         - rgb_array: Return an numpy.ndarray with shape (x, y, 3),           representing RGB values for an x-by-y pixel image, suitable           for turning into a video.         - ansi: Return a string (str) or StringIO.StringIO containing a           terminal-style text representation. The text can include newlines           and ANSI escape sequences (e.g. for colors).          Note:             Make sure that your class's metadata 'render.modes' key includes               the list of supported modes. It's recommended to call super()               in implementations to use the functionality of this method.          Args:             mode (str): the mode to render with          Example:          class MyEnv(Env):             metadata = {'render.modes': ['human', 'rgb_array']}              def render(self, mode='human'):                 if mode == 'rgb_array':                     return np.array(...) # return RGB frame suitable for video                 elif mode == 'human':                     ... # pop up a window and render                 else:                     super(MyEnv, self).render(mode=mode) # just raise an exception Override close in your subclass to perform any necessary cleanup.          Environments will automatically close() themselves when         garbage collected or when the program exits. Sets the seed for this env's random number generator(s).          Note:             Some environments use multiple pseudorandom number generators.             We want to capture all such seeds used in order to ensure that             there aren't accidental correlations between multiple generators.          Returns:             list<bigint>: Returns the list of seeds used in this env's random               number generators. The first value in the list should be the               "main" seed, or the value which a reproducer should pass to               'seed'. Often, the main seed equals the provided 'seed', but               this won't be true if seed=None, for example. Completely unwrap this env.          Returns:             gym.Env: The base non-wrapped gym.Env instance Support with-statement for the environment. Support with-statement for the environment. # propagate exception # th := theta # since the euler method used for dt = 0.001 # for rendering