Created test case for A3C

"""Interface for reinforcement learning."""

from deepchem.rl.a3c import A3C

class Environment(object):

  """An environment in which an actor performs actions to accomplish a task.

  "action" argument passed to the environment is an integer, giving the index of the action to perform.

  """

  def __init__(self, env, policy, max_rollout_length=20, discount_factor=0.99, value_weight=0.25, entropy_weight=0.01):

  def __init__(self, env, policy, max_rollout_length=20, discount_factor=0.99, value_weight=1.0, entropy_weight=0.01):

    """Create an object for optimizing a policy.

    Parameters

    session = self.a3c._session

    states = []

    actions = []

    scores = []

    rewards = []

    for i in range(self.a3c.max_rollout_length):

      if self.env.terminated:

        break

      action = np.random.choice(np.arange(n_actions), p=probabilities[0])

      actions.append(np.zeros(n_actions))

      actions[i][action] = 1.0

      scores.append(self.env.step(action))

      rewards.append(self.env.step(action))

    if not self.env.terminated:

      # Add an estimate of the reward for the rest of the episode.

      feed_dict = {self.features.out_tensor: np.expand_dims(self.env.state, axis=0)}

      scores[-1] += session.run(self.value.out_tensor, feed_dict)

    for j in range(len(scores)-1, 0, -1):

      scores[j-1] += self.a3c.discount_factor*scores[j]

      rewards[-1] += self.a3c.discount_factor*session.run(self.value.out_tensor, feed_dict)

    for j in range(len(rewards)-1, 0, -1):

      rewards[j-1] += self.a3c.discount_factor*rewards[j]

    if self.env.terminated:

      self.env.reset()

    return np.array(states), np.array(actions), np.array(scores).reshape((len(scores), 1))

    return np.array(states), np.array(actions), np.array(rewards).reshape((len(rewards), 1))

import deepchem as dc

from deepchem.models.tensorgraph.layers import Reshape, Dense, SoftMax

import numpy as np

import tensorflow as tf

import unittest

class TestA3C(unittest.TestCase):

  def test_roulette(self):

    """Test training a policy for the roulette environment."""

    # This is modeled after the Roulette-v0 environment from OpenAI Gym.

    # The player can bet on any number from 0 to 36, or walk away (which ends the

    # game).  The average reward for any bet is slightly negative, so the best

    # strategy is to walk away.

    class RouletteEnvironment(dc.rl.Environment):

      def __init__(self):

        super().__init__([1], 38)

        self._state = np.array([0])

      def step(self, action):

        if action == 37:

          self._terminated = True # Walk away.

          return 0.0

        wheel = np.random.randint(37)

        if wheel == 0:

          if action == 0:

            return 35.0

          return -1.0

        if action != 0 and wheel%2 == action%2:

          return 1.0

        return -1.0

      def reset(self):

        self._terminated = False

    env = RouletteEnvironment()

    # This policy just learns a constant probability for each action, and a constant for the value.

    class TestPolicy(dc.rl.Policy):

      def create_layers(self, features, **kwargs):

        action = Dense(in_layers=[features], out_channels=env.n_actions)

        output = SoftMax(in_layers=[Reshape(in_layers=[action], shape=(-1, env.n_actions))])

        value = Dense(in_layers=[features], out_channels=1)

        return {'action_prob':output, 'value':value}

    # Optimize it.

    a3c = dc.rl.A3C(env, TestPolicy(), value_weight=100.0)

    a3c.optimizer = dc.models.tensorgraph.TFWrapper(tf.train.AdamOptimizer, learning_rate=0.1)

    a3c.fit(100000)

    # It should have learned that the expected value is very close to zero, and that the best

    # action is to walk away.

    action_prob, value = a3c.predict([0])

    assert -0.5 < value[0] < 0.5

    assert action_prob.argmax() == 37

    assert a3c.select_action([0], deterministic=True) == 37