malyvsen / wombat Goto Github PK

import gym # OpenAI gym supported out of the box
import wombat
env = gym.make('CartPole-v0')
agent = wombat.agents.RandomDiscrete(env.action_space.n)
wombat.run(agent, env, num_episodes=4, per_step=env.render)

class CyclicAgent: # no base class needed
   def __init__(self, num_possible_actions):
      self.num_possible_actions = num_possible_actions
   def act(self, steps): # should return selected action
      return len(steps) % self.num_possible_actions
   def train(self, steps): # should return mean training loss (eg. for prioritized experience replay)
      return 0

cyclic_agent = CyclicAgent(env.action_space.n)
wombat.run(cyclic_agent, env, num_episodes=4, per_step=env.render)

episode = wombat.Episode() # will record all steps
for step in episode.run(agent, env):
   print(f'Action {step.action} resulted in reward of {step.reward}')
print(f'Episode finished, total reward: {episode.total_reward()}')
agent.train(episode.steps) # episode can be used with wombat, just like that