diff --git a/test/test_utils.py b/test/test_utils.py index b0b5325..7c07d49 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -1,10 +1,16 @@ from rl_portfolio_management.data.utils import random_shift, normalize, scale_to_start -from rl_portfolio_management.util import sharpe, MDD +from rl_portfolio_management.util import sharpe, MDD, softmax import pandas as pd import numpy as np +def test_softmax(): + x = np.random.random((20, 20)) + y = softmax(x) + np.testing.assert_almost_equal(y.sum(), 1) + + def test_maxdrawdown(): assert MDD(np.array([0, 0, 0, 0, 1, 2, 3])) == 0 assert MDD(np.array([0, 0, 0, 0, 1, 2, 1])) == -1 diff --git a/test/test_wrappers.py b/test/test_wrappers.py new file mode 100644 index 0000000..460efd5 --- /dev/null +++ b/test/test_wrappers.py @@ -0,0 +1,24 @@ +import gym + +from rl_portfolio_management.environments import PortfolioEnv, env_specs +from rl_portfolio_management.wrappers import ConcatStates, SoftmaxActions + + +def test_concat(): + env = gym.make("CryptoPortfolioEIIE-v0") + env = ConcatStates(env) + obs = env.reset() + assert len(obs.shape) == 3 + action = env.action_space.sample() + action /= action.sum() + obs, rew, done, info = env.step(action) + assert len(obs.shape) == 3 + + +def test_softmax(): + env = gym.make("CryptoPortfolioEIIE-v0") + env = SoftmaxActions(env) + obs = env.reset() + # should be no problem with actions that don't sum to one + action = env.action_space.sample() * 100 + obs, rew, done, info = env.step(action)