Loading deepchem/rl/a3c.py +2 −2 Original line number Diff line number Diff line Loading @@ -96,7 +96,7 @@ class A3C(object): env: Environment the Environment to interact with policy: Policy the Policy to optimize. Its create_layers() method must return a map containing the the Policy to optimize. Its create_layers() method must return a dict containing the keys 'action_prob' and 'value', corresponding to the action probabilities and value estimate max_rollout_length: int the maximum length of rollouts to generate Loading Loading @@ -409,7 +409,7 @@ class _Worker(object): if self.env.terminated: self.env.reset() self.rnn_states = self.graph.rnn_zero_states return states, actions, np.array(rewards), np.array(values) return states, actions, np.array(rewards, dtype=np.float32), np.array(values, dtype=np.float32) def process_rollout(self, states, actions, rewards, values, initial_rnn_states, step_count): Loading deepchem/rl/mcts.py +52 −46 Original line number Diff line number Diff line Loading @@ -87,7 +87,7 @@ class MCTS(object): env: Environment the Environment to interact with policy: Policy the Policy to optimize. Its create_layers() method must return a map containing the the Policy to optimize. Its create_layers() method must return a dict containing the keys 'action_prob' and 'value', corresponding to the action probabilities and value estimate max_search_depth: int the maximum depth of the tree search, measured in steps Loading Loading @@ -218,14 +218,17 @@ class MCTS(object): # Run the algorithm. for iteration in range(iterations): buffer = self._run_episodes(steps_per_iteration, temperature, saver, adapt_puct) buffer = self._run_episodes(steps_per_iteration, temperature, saver, adapt_puct) self._optimize_policy(buffer, epochs_per_iteration) # Save a file checkpoint. self._checkpoint_index += 1 saver.save( self._graph.session, self._graph.save_file, global_step=self._checkpoint_index) self._graph.session, self._graph.save_file, global_step=self._checkpoint_index) def predict(self, state): """Compute the policy's output predictions for a state. Loading @@ -247,9 +250,7 @@ class MCTS(object): results = self._graph.session.run(tensors, feed_dict=feed_dict) return results[:2] def select_action(self, state, deterministic=False): def select_action(self, state, deterministic=False): """Select an action to perform based on the environment's state. Parameters Loading @@ -268,7 +269,8 @@ class MCTS(object): state = [state] with self._graph._get_tf("Graph").as_default(): feed_dict = self._create_feed_dict(state) probabilities = self._graph.session.run(self._pred_prob, feed_dict=feed_dict) probabilities = self._graph.session.run( self._pred_prob, feed_dict=feed_dict) if deterministic: return probabilities.argmax() else: Loading Loading @@ -309,12 +311,13 @@ class MCTS(object): self._env.reset() root = TreeSearchNode(0.0) else: # root = root.children[action] root = TreeSearchNode(0.0) root = root.children[action] if time.time() > self._checkpoint_time: self._checkpoint_index += 1 saver.save( self._graph.session, self._graph.save_file, global_step=self._checkpoint_index) self._graph.session, self._graph.save_file, global_step=self._checkpoint_index) self._checkpoint_time = time.time() return buffer Loading @@ -332,10 +335,12 @@ class MCTS(object): for i, f in enumerate(self._features): feed_dict[f] = np.stack(buffer[j][0][i] for j in indices) feed_dict[self._search_prob] = np.stack(buffer[j][1] for j in indices) feed_dict[self._search_value] = np.array([buffer[j][2] for j in indices]) feed_dict[self._search_value] = np.array( [buffer[j][2] for j in indices]) yield feed_dict loss = self._graph.fit_generator(generate_batches(), checkpoint_interval=0) loss = self._graph.fit_generator( generate_batches(), checkpoint_interval=0) def _do_tree_search(self, root, temperature, adapt_puct): """Perform the tree search for a state.""" Loading @@ -343,9 +348,7 @@ class MCTS(object): for i in range(self.n_search_episodes): env = copy.deepcopy(self._env) trace = [] self._create_trace(env, root, trace) self._record_trace_rewards(trace) self._create_trace(env, root, 1) # Compute the final probabilities and expected reward. Loading @@ -353,27 +356,28 @@ class MCTS(object): prob /= np.sum(prob) reward = np.sum(p * c.mean_reward for p, c in zip(prob, root.children)) if adapt_puct: scale = np.sum([p*np.abs(c.mean_reward) for p, c in zip(prob, root.children)]) scale = np.sum( [p * np.abs(c.mean_reward) for p, c in zip(prob, root.children)]) self._puct_scale = 0.99 * self._puct_scale + 0.01 * scale return prob, reward def _create_trace(self, env, node, trace): def _create_trace(self, env, node, depth): """Create one trace as part of the tree search.""" trace.append(node) node.count += 1 if env.terminated: # Mark this node as terminal node.children = None node.value = 0.0 return return 0.0 if node.children is not None and len(node.children) == 0: # Expand this node. prob_pred, value = self.predict(env.state) node.value = float(value) node.children = [TreeSearchNode(p) for p in prob_pred[0]] if len(trace) == self.max_search_depth: return if depth == self.max_search_depth: reward = 0.0 future_rewards = node.value else: # Select the next action to perform. total_counts = sum(c.count for c in node.children) Loading @@ -381,22 +385,24 @@ class MCTS(object): score = [c.prior_prob for c in node.children] else: scale = self._puct_scale * np.sqrt(total_counts) score = [c.mean_reward + scale*c.prior_prob/(1+c.count) for c in node.children] score = [ c.mean_reward + scale * c.prior_prob / (1 + c.count) for c in node.children ] action = np.argmax(score) next_node = node.children[action] next_node.reward = env.step(action) reward = env.step(action) # Recursively build the tree. self._create_trace(env, next_node, trace) future_rewards = self._create_trace(env, next_node, depth + 1) # Update statistics for this node. def _record_trace_rewards(self, trace): """Record the rewards received during a trace.""" value = trace[-1].value for node in reversed(trace): value = node.reward + self.discount_factor*value node.total_reward += value future_rewards = reward + self.discount_factor * future_rewards node.total_reward += future_rewards node.mean_reward = node.total_reward / node.count return future_rewards class TreeSearchNode(object): Loading deepchem/rl/ppo.py +1 −1 Original line number Diff line number Diff line Loading @@ -107,7 +107,7 @@ class PPO(object): env: Environment the Environment to interact with policy: Policy the Policy to optimize. Its create_layers() method must return a map containing the the Policy to optimize. Its create_layers() method must return a dict containing the keys 'action_prob' and 'value', corresponding to the action probabilities and value estimate max_rollout_length: int the maximum length of rollouts to generate Loading deepchem/rl/tests/test_mcts.py +1 −1 Original line number Diff line number Diff line Loading @@ -63,7 +63,7 @@ class TestMCTS(unittest.TestCase): max_search_depth=5, n_search_episodes=200, optimizer=Adam(learning_rate=0.005)) mcts.fit(3, steps_per_iteration=50, epochs_per_iteration=50, puct_scale=1.0) mcts.fit(10, steps_per_iteration=50, epochs_per_iteration=50) # It should have learned that the expected value is very close to zero, and that the best # action is to walk away. Loading Loading
deepchem/rl/a3c.py +2 −2 Original line number Diff line number Diff line Loading @@ -96,7 +96,7 @@ class A3C(object): env: Environment the Environment to interact with policy: Policy the Policy to optimize. Its create_layers() method must return a map containing the the Policy to optimize. Its create_layers() method must return a dict containing the keys 'action_prob' and 'value', corresponding to the action probabilities and value estimate max_rollout_length: int the maximum length of rollouts to generate Loading Loading @@ -409,7 +409,7 @@ class _Worker(object): if self.env.terminated: self.env.reset() self.rnn_states = self.graph.rnn_zero_states return states, actions, np.array(rewards), np.array(values) return states, actions, np.array(rewards, dtype=np.float32), np.array(values, dtype=np.float32) def process_rollout(self, states, actions, rewards, values, initial_rnn_states, step_count): Loading
deepchem/rl/mcts.py +52 −46 Original line number Diff line number Diff line Loading @@ -87,7 +87,7 @@ class MCTS(object): env: Environment the Environment to interact with policy: Policy the Policy to optimize. Its create_layers() method must return a map containing the the Policy to optimize. Its create_layers() method must return a dict containing the keys 'action_prob' and 'value', corresponding to the action probabilities and value estimate max_search_depth: int the maximum depth of the tree search, measured in steps Loading Loading @@ -218,14 +218,17 @@ class MCTS(object): # Run the algorithm. for iteration in range(iterations): buffer = self._run_episodes(steps_per_iteration, temperature, saver, adapt_puct) buffer = self._run_episodes(steps_per_iteration, temperature, saver, adapt_puct) self._optimize_policy(buffer, epochs_per_iteration) # Save a file checkpoint. self._checkpoint_index += 1 saver.save( self._graph.session, self._graph.save_file, global_step=self._checkpoint_index) self._graph.session, self._graph.save_file, global_step=self._checkpoint_index) def predict(self, state): """Compute the policy's output predictions for a state. Loading @@ -247,9 +250,7 @@ class MCTS(object): results = self._graph.session.run(tensors, feed_dict=feed_dict) return results[:2] def select_action(self, state, deterministic=False): def select_action(self, state, deterministic=False): """Select an action to perform based on the environment's state. Parameters Loading @@ -268,7 +269,8 @@ class MCTS(object): state = [state] with self._graph._get_tf("Graph").as_default(): feed_dict = self._create_feed_dict(state) probabilities = self._graph.session.run(self._pred_prob, feed_dict=feed_dict) probabilities = self._graph.session.run( self._pred_prob, feed_dict=feed_dict) if deterministic: return probabilities.argmax() else: Loading Loading @@ -309,12 +311,13 @@ class MCTS(object): self._env.reset() root = TreeSearchNode(0.0) else: # root = root.children[action] root = TreeSearchNode(0.0) root = root.children[action] if time.time() > self._checkpoint_time: self._checkpoint_index += 1 saver.save( self._graph.session, self._graph.save_file, global_step=self._checkpoint_index) self._graph.session, self._graph.save_file, global_step=self._checkpoint_index) self._checkpoint_time = time.time() return buffer Loading @@ -332,10 +335,12 @@ class MCTS(object): for i, f in enumerate(self._features): feed_dict[f] = np.stack(buffer[j][0][i] for j in indices) feed_dict[self._search_prob] = np.stack(buffer[j][1] for j in indices) feed_dict[self._search_value] = np.array([buffer[j][2] for j in indices]) feed_dict[self._search_value] = np.array( [buffer[j][2] for j in indices]) yield feed_dict loss = self._graph.fit_generator(generate_batches(), checkpoint_interval=0) loss = self._graph.fit_generator( generate_batches(), checkpoint_interval=0) def _do_tree_search(self, root, temperature, adapt_puct): """Perform the tree search for a state.""" Loading @@ -343,9 +348,7 @@ class MCTS(object): for i in range(self.n_search_episodes): env = copy.deepcopy(self._env) trace = [] self._create_trace(env, root, trace) self._record_trace_rewards(trace) self._create_trace(env, root, 1) # Compute the final probabilities and expected reward. Loading @@ -353,27 +356,28 @@ class MCTS(object): prob /= np.sum(prob) reward = np.sum(p * c.mean_reward for p, c in zip(prob, root.children)) if adapt_puct: scale = np.sum([p*np.abs(c.mean_reward) for p, c in zip(prob, root.children)]) scale = np.sum( [p * np.abs(c.mean_reward) for p, c in zip(prob, root.children)]) self._puct_scale = 0.99 * self._puct_scale + 0.01 * scale return prob, reward def _create_trace(self, env, node, trace): def _create_trace(self, env, node, depth): """Create one trace as part of the tree search.""" trace.append(node) node.count += 1 if env.terminated: # Mark this node as terminal node.children = None node.value = 0.0 return return 0.0 if node.children is not None and len(node.children) == 0: # Expand this node. prob_pred, value = self.predict(env.state) node.value = float(value) node.children = [TreeSearchNode(p) for p in prob_pred[0]] if len(trace) == self.max_search_depth: return if depth == self.max_search_depth: reward = 0.0 future_rewards = node.value else: # Select the next action to perform. total_counts = sum(c.count for c in node.children) Loading @@ -381,22 +385,24 @@ class MCTS(object): score = [c.prior_prob for c in node.children] else: scale = self._puct_scale * np.sqrt(total_counts) score = [c.mean_reward + scale*c.prior_prob/(1+c.count) for c in node.children] score = [ c.mean_reward + scale * c.prior_prob / (1 + c.count) for c in node.children ] action = np.argmax(score) next_node = node.children[action] next_node.reward = env.step(action) reward = env.step(action) # Recursively build the tree. self._create_trace(env, next_node, trace) future_rewards = self._create_trace(env, next_node, depth + 1) # Update statistics for this node. def _record_trace_rewards(self, trace): """Record the rewards received during a trace.""" value = trace[-1].value for node in reversed(trace): value = node.reward + self.discount_factor*value node.total_reward += value future_rewards = reward + self.discount_factor * future_rewards node.total_reward += future_rewards node.mean_reward = node.total_reward / node.count return future_rewards class TreeSearchNode(object): Loading
deepchem/rl/ppo.py +1 −1 Original line number Diff line number Diff line Loading @@ -107,7 +107,7 @@ class PPO(object): env: Environment the Environment to interact with policy: Policy the Policy to optimize. Its create_layers() method must return a map containing the the Policy to optimize. Its create_layers() method must return a dict containing the keys 'action_prob' and 'value', corresponding to the action probabilities and value estimate max_rollout_length: int the maximum length of rollouts to generate Loading
deepchem/rl/tests/test_mcts.py +1 −1 Original line number Diff line number Diff line Loading @@ -63,7 +63,7 @@ class TestMCTS(unittest.TestCase): max_search_depth=5, n_search_episodes=200, optimizer=Adam(learning_rate=0.005)) mcts.fit(3, steps_per_iteration=50, epochs_per_iteration=50, puct_scale=1.0) mcts.fit(10, steps_per_iteration=50, epochs_per_iteration=50) # It should have learned that the expected value is very close to zero, and that the best # action is to walk away. Loading
