SineTradeEnv-v0

trade_env.py

# trade_env.py
import numpy as np
import gym

class SineEnv(gym.Env):
    """
    An Env that trades for n steps and returns rewards and observations by following
    a sine curve
    """
    def __init__(self, n=200, offset=3, amplitude=0.6, wavelength=0.1):
        
        self.n = n
        self.offset = offset
        self.amplitude = amplitude
        self.wavelength = wavelength
        
        self.num_actions = 3
        self.series = self.offset + self.amplitude * np.sin(self.wavelength * np.arange(self.n))
        self.done = False
    
    def step(self, action):
        """
        The step function takes action and returns the next obs, reward, done and info
        """
        self.index += 1
        obs =  self.series[self.index]
        step_return = np.log(self.series[self.index]/self.series[self.index-1])
        reward = (action - 1) * step_return
        
        if self.index >= (self.n-1):
            self.done = True
            
        info = {
            'return': step_return
        }
        return obs, reward, self.done, info
        
    def reset(self):
        self.index = 0
        self.done = False
        obs = self.series[self.index]
        return obs