# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/models/layers/models.layers.ou_noise.ipynb (unless otherwise specified).

__all__ = ['OUNoise']

# Cell
import torch

# Cell
class OUNoise(object):
    """
    Ornstein-Uhlenbeck Noise
    """
    def __init__(self,
                 embedded_action_size,
                 ou_mu,
                 ou_theta,
                 ou_sigma,
                 ou_epsilon):
        """
        Initialize OUNoise
        """
        self.embedded_action_size = embedded_action_size
        self.ou_mu = ou_mu
        self.ou_theta = ou_theta
        self.ou_sigma = ou_sigma
        self.ou_epsilon = ou_epsilon
        self.ou_state = None
        self.reset()

    def reset(self):
        """
        Reset the OU process state
        """
        self.ou_state = torch.ones(self.embedded_action_size) * self.ou_mu

    def evolve_state(self):
        """
        Evolve the OU process state
        """
        self.ou_state += self.ou_theta * (self.ou_mu - self.ou_state) \
            + self.ou_sigma * torch.randn(self.embedded_action_size)

    def get_ou_noise(self):
        """
        Get the OU noise for one action
        :return OU noise
        """
        self.evolve_state()
        return self.ou_state.clone()