forked from Kaggle/docker-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_tensorforce.py
More file actions
30 lines (25 loc) · 1.08 KB
/
test_tensorforce.py
File metadata and controls
30 lines (25 loc) · 1.08 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
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)