docker-python/tests/test_tensorforce.py at master · Santhosh-Suresh/docker-python · GitHub

30 lines (25 loc) · 1.08 KB

import unittest

from tensorforce import Agent, Environment

class TestTensorforce(unittest.TestCase):

# based on https://github.com/tensorforce/tensorforce/tree/master#quickstart-example-code.

def test_quickstart(self):

environment = Environment.create(

environment='gym', level='CartPole', max_episode_timesteps=500

)

agent = Agent.create(

agent='tensorforce',

environment=environment, # alternatively: states, actions, (max_episode_timesteps)

memory=1000,

update=dict(unit='timesteps', batch_size=32),

optimizer=dict(type='adam', learning_rate=3e-4),

policy=dict(network='auto'),

objective='policy_gradient',

reward_estimation=dict(horizon=1)

)

# Train for a single episode.

states = environment.reset()

actions = agent.act(states=states)

states, terminal, reward = environment.execute(actions=actions)

self.assertEqual(4, len(states))

self.assertFalse(terminal)

self.assertEqual(1, reward)

Expand file tree