<filename>rlkit/envs/obs_wrappers.py
import importlib
from typing import Any, Dict
import numpy as np
import gym
from gym import Env
from gym.spaces import Box
from collections import deque

from torch import matmul
from rlkit.core import logger

from rlkit.core.serializable import Serializable
from rlkit.envs.wrappers import ProxyEnv

EPS = np.finfo(np.float32).eps.item()


class Process:
    
    def __call__(self, size: int = None, **kwargs) -> Any:
        pass

    def reset(self) -> None:
        pass


class Bernoulli(Process):

    def __init__(self, p: float = 0.5, low: float = 0.0, high: float = 1.0):
        self.p = p
        self.low = low
        self.high = high

    def __call__(self, size: int = None, **kwargs) -> np.ndarray:
        v = int(np.random.random(size=size) < self.p)
        v[v == 0] = -np.inf
        v[v == 1] = np.high
        v[np.isinf(v)] == np.low
        return v


class Randomwalk(Process):

    def __init__(self, p: float = 0.5):
        self.p = p
        self.x = None

    def __call__(self, size: int = None, **kwargs) -> np.ndarray:
        if self.x is None:
            self.x = np.zeros(shape=size if size else 1)
            return self.x
        v = np.random.random(size=size) < self.p
        v = np.where(v, 1, -1)
        self.x += v
        return self.x

    def reset(self) -> None:
        self.x = None


class Gaussian(Process):
    
    def __init__(self, mean: float = 0.0, std: float = 1.0):
        self.mean = mean
        self.std = std

    def __call__(self, size: int = None, **kwargs) -> np.ndarray:
        return self.mean + np.random.normal(size=size) * self.std


class Uniform(Process):

    def __init__(self, low: float = 0.0, high: float = 1.0):
        self.low = low
        self.high =high

    def __call__(self, size: int = None, **kwargs) -> np.ndarray:
        return np.random.uniform(low=self.low, high=self.high, size=size)
        

class Wiener(Process):

    def __init__(self, std: float = 1.0, mu: float = 0.0):
        self.std = std
        self.mu = mu
        self.w = None

    def __call__(self, size: int = None, **kwargs) -> np.ndarray:
        if self.w is None:
            self.w = np.zeros(shape=size if size else 1)
            return self.w
        self.w = self.w + self.mu + np.random.normal(loc=0.0, scale=self.std, size=size)
        return self.w

    def reset(self) -> None:
        self.w = None


class Vagram(Process):
    """
    This design is from the paper <NAME>, <NAME>, et. al.,
    "Value Gradient Weighted Model-Based Reinforcement Learning"
    https://openreview.net/forum?id=4-D6CZkRXxI
    """

    def __init__(self, s_std: float = 0.1, m_std: float = 10.0):
        self.s_std = s_std
        self.m_std = m_std
        self.s0 = None
        self.A = None
        self.l = None
        self.s = None

    def __call__(self, size: int = None, **kwargs) -> np.ndarray:
        if self.s0 is None:
            assert size is not None
            self.s0 = np.random.normal(scale=self.s_std, size=size)
            self.s = self.s0.copy()
            self.l = np.random.normal(size=[size, size])
            self.A = np.random.normal(scale=self.m_std, size=[size, size]) ** 2
            return self.s
        self.s = np.matmul(self.l, self.s)
        self.s += 0.1 * np.sin(np.matmul(self.A, self.s))
        self.s = np.where(np.abs(self.s) > 20.0, self.s0, self.s)
        return self.s

    def reset(self) -> None:
        self.s0 = None
        self.A = None
        self.s = None
        self.l = None


class ExtraObsEnv(ProxyEnv, Serializable):
    """
    Add extra dimensions to observation spaces
    """

    def __init__(
        self, 
        env: Env,
        extra_num: int, # number of extra dimensions
        proc: str = 'gaussian', # process to generate extra observations
        proc_args: Dict[str, Any] = None
    ):
        self._wrapped_env = env
        self._serializable_initialized = False
        Serializable.quick_init(self, locals())
        ProxyEnv.__init__(self, env)

        self.extra_num = extra_num
        obs: Box = self.observation_space
        assert isinstance(obs, Box)
        self.observation_space = Box(
            low=np.concatenate([obs.low, [-np.inf] * extra_num]),
            high=np.concatenate([obs.high, [np.inf] * extra_num]),
            dtype=np.float32
        )
        proc = proc[0].upper() + proc[1:]
        this_module = importlib.import_module(name='.obs_wrappers', package='rlkit.envs')
        self.proc: Process = getattr(this_module, proc)(**proc_args)
        print(f'Use noise wrapper: {proc} | noise kwargs: {proc_args}')

    def _extend_obs(self, obs: np.ndarray) -> np.ndarray:
        assert obs.ndim == 1
        noise = self.proc(size=self.extra_num)
        new_obs = np.concatenate([obs, noise])
        return new_obs

    def reset(self, **kwargs):
        self.proc.reset()
        obs = self._wrapped_env.reset(**kwargs)
        return self._extend_obs(obs)

    def step(self, action):
        obs, r, done, info = self._wrapped_env.step(action)
        obs = self._extend_obs(obs)
        return obs, r, done, info

    def log_statistics(self, *args, **kwargs) -> Dict:
        if hasattr(self._wrapped_env, "log_statistics"):
            return self._wrapped_env.log_statistics(*args, **kwargs)
        else:
            return {}

    def log_new_ant_multi_statistics(self, paths, epoch, log_dir) -> Dict:
        if hasattr(self._wrapped_env, "log_new_ant_multi_statistics"):
            return self._wrapped_env.log_new_ant_multi_statistics(paths, epoch, log_dir)
        else:
            return {}
