diff --git a/examples/visualization_demo.py b/examples/visualization_demo.py index 46549db..0af450b 100644 --- a/examples/visualization_demo.py +++ b/examples/visualization_demo.py @@ -52,7 +52,7 @@ def convert_tree_to_graph( chex.assert_rank(tree.node_values, 2) batch_size = tree.node_values.shape[0] if action_labels is None: - action_labels = range(tree.num_actions) + action_labels = range(tree.num_actions) # pyrefly: ignore[bad-assignment] elif len(action_labels) != tree.num_actions: raise ValueError( f"action_labels {action_labels} has the wrong number of actions " @@ -60,6 +60,7 @@ def convert_tree_to_graph( f"Expecting {tree.num_actions}.") def node_to_str(node_i, reward=0, discount=1): + # pyrefly: ignore[bad-index] return (f"{node_i}\n" f"Reward: {reward:.2f}\n" f"Discount: {discount:.2f}\n" @@ -68,7 +69,8 @@ def node_to_str(node_i, reward=0, discount=1): def edge_to_str(node_i, a_i): node_index = jnp.full([batch_size], node_i) - probs = jax.nn.softmax(tree.children_prior_logits[batch_index, node_i]) + probs = jax.nn.softmax(tree.children_prior_logits[batch_index, node_i]) # pyrefly: ignore[bad-index] + # pyrefly: ignore[unsupported-operation] return (f"{action_labels[a_i]}\n" f"Q: {tree.qvalues(node_index)[batch_index, a_i]:.2f}\n" # pytype: disable=unsupported-operands # always-use-return-annotations f"p: {probs[a_i]:.2f}\n") @@ -81,14 +83,14 @@ def edge_to_str(node_i, a_i): for node_i in range(tree.num_simulations): for a_i in range(tree.num_actions): # Index of children, or -1 if not expanded - children_i = tree.children_index[batch_index, node_i, a_i] + children_i = tree.children_index[batch_index, node_i, a_i] # pyrefly: ignore[bad-index] if children_i >= 0: graph.add_node( children_i, label=node_to_str( node_i=children_i, - reward=tree.children_rewards[batch_index, node_i, a_i], - discount=tree.children_discounts[batch_index, node_i, a_i]), + reward=tree.children_rewards[batch_index, node_i, a_i], # pyrefly: ignore[bad-index] + discount=tree.children_discounts[batch_index, node_i, a_i]), # pyrefly: ignore[bad-index] color="red") graph.add_edge(node_i, children_i, label=edge_to_str(node_i, a_i)) @@ -166,8 +168,8 @@ def _make_batched_env_model( root_state = 0 root = mctx.RootFnOutput( prior_logits=jnp.full([batch_size, num_actions], - prior_logits[root_state]), - value=jnp.full([batch_size], values[root_state]), + prior_logits[root_state]), # pyrefly: ignore[bad-index] + value=jnp.full([batch_size], values[root_state]), # pyrefly: ignore[bad-index] # The embedding will hold the state index. embedding=jnp.zeros([batch_size], dtype=jnp.int32), ) @@ -177,11 +179,11 @@ def recurrent_fn(params, rng_key, action, embedding): chex.assert_shape(action, [batch_size]) chex.assert_shape(embedding, [batch_size]) recurrent_fn_output = mctx.RecurrentFnOutput( - reward=rewards[embedding, action], - discount=discounts[embedding, action], - prior_logits=prior_logits[embedding], - value=values[embedding]) - next_embedding = transition_matrix[embedding, action] + reward=rewards[embedding, action], # pyrefly: ignore[bad-index] + discount=discounts[embedding, action], # pyrefly: ignore[bad-index] + prior_logits=prior_logits[embedding], # pyrefly: ignore[bad-index] + value=values[embedding]) # pyrefly: ignore[bad-index] + next_embedding = transition_matrix[embedding, action] # pyrefly: ignore[bad-index] return recurrent_fn_output, next_embedding return root, recurrent_fn diff --git a/mctx/_src/action_selection.py b/mctx/_src/action_selection.py index f09aca9..977cd35 100644 --- a/mctx/_src/action_selection.py +++ b/mctx/_src/action_selection.py @@ -71,14 +71,16 @@ def muzero_action_selection( Returns: action: the action selected from the given node. """ - visit_counts = tree.children_visits[node_index] - node_visit = tree.node_visits[node_index] + visit_counts = tree.children_visits[node_index] # pyrefly: ignore[bad-index] + node_visit = tree.node_visits[node_index] # pyrefly: ignore[bad-index] pb_c = pb_c_init + jnp.log((node_visit + pb_c_base + 1.) / pb_c_base) + # pyrefly: ignore[bad-index] prior_logits = tree.children_prior_logits[node_index] prior_probs = jax.nn.softmax(prior_logits) policy_score = jnp.sqrt(node_visit) * pb_c * prior_probs / (visit_counts + 1) chex.assert_shape([node_index, node_visit], ()) chex.assert_equal_shape([prior_probs, visit_counts, policy_score]) + # pyrefly: ignore[bad-argument-type] value_score = qtransform(tree, node_index) # Add tiny bit of randomness for tie break @@ -129,9 +131,12 @@ def gumbel_muzero_root_action_selection( """ del rng_key chex.assert_shape([node_index], ()) + # pyrefly: ignore[bad-index] visit_counts = tree.children_visits[node_index] + # pyrefly: ignore[bad-index] prior_logits = tree.children_prior_logits[node_index] chex.assert_equal_shape([visit_counts, prior_logits]) + # pyrefly: ignore[bad-argument-type] completed_qvalues = qtransform(tree, node_index) table = jnp.array(seq_halving.get_table_of_considered_visits( @@ -180,9 +185,11 @@ def gumbel_muzero_interior_action_selection( """ del rng_key, depth chex.assert_shape([node_index], ()) - visit_counts = tree.children_visits[node_index] + visit_counts = tree.children_visits[node_index] # pyrefly: ignore[bad-index] + # pyrefly: ignore[bad-index] prior_logits = tree.children_prior_logits[node_index] chex.assert_equal_shape([visit_counts, prior_logits]) + # pyrefly: ignore[bad-argument-type] completed_qvalues = qtransform(tree, node_index) # The `prior_logits + completed_qvalues` provide an improved policy, diff --git a/mctx/_src/policies.py b/mctx/_src/policies.py index f5f6538..3df7d7f 100644 --- a/mctx/_src/policies.py +++ b/mctx/_src/policies.py @@ -86,7 +86,7 @@ def muzero_policy( jax.nn.softmax(root.prior_logits), dirichlet_fraction=dirichlet_fraction, dirichlet_alpha=dirichlet_alpha)) - root = root.replace( + root = root.replace( # pyrefly: ignore[missing-attribute] prior_logits=_mask_invalid_actions(noisy_logits, invalid_actions)) # Running the search. @@ -176,7 +176,7 @@ def gumbel_muzero_policy( search tree. """ # Masking invalid actions. - root = root.replace( + root = root.replace( # pyrefly: ignore[missing-attribute] prior_logits=_mask_invalid_actions(root.prior_logits, invalid_actions)) # Generating Gumbel. @@ -301,7 +301,7 @@ def stochastic_muzero_policy( dirichlet_fraction=dirichlet_fraction, dirichlet_alpha=dirichlet_alpha)) - root = root.replace( + root = root.replace( # pyrefly: ignore[missing-attribute] prior_logits=_mask_invalid_actions(noisy_logits, invalid_actions)) # construct a dummy afterstate embedding @@ -509,7 +509,7 @@ def _take_slice(x): else: raise ValueError(f'Unknown mode: {mode}.') - return tree.replace( + return tree.replace( # pyrefly: ignore[missing-attribute] children_index=_take_slice(tree.children_index), children_prior_logits=_take_slice(tree.children_prior_logits), children_visits=_take_slice(tree.children_visits), @@ -531,7 +531,8 @@ def _chance_node_selection_fn( tree: search.Tree, node_index: chex.Array, ) -> chex.Array: - num_chance = tree.children_visits[node_index] + num_chance = tree.children_visits[node_index] # pyrefly: ignore[bad-index] + # pyrefly: ignore[bad-index] chance_logits = tree.children_prior_logits[node_index] prob_chance = jax.nn.softmax(chance_logits) argmax_chance = jnp.argmax(prob_chance / (num_chance + 1), axis=-1).astype( diff --git a/mctx/_src/qtransforms.py b/mctx/_src/qtransforms.py index 08ee9b8..1280ab5 100644 --- a/mctx/_src/qtransforms.py +++ b/mctx/_src/qtransforms.py @@ -44,8 +44,9 @@ def qtransform_by_min_max( """ chex.assert_shape(node_index, ()) qvalues = tree.qvalues(node_index) - visit_counts = tree.children_visits[node_index] + visit_counts = tree.children_visits[node_index] # pyrefly: ignore[bad-index] value_score = jnp.where(visit_counts > 0, qvalues, min_value) + # pyrefly: ignore[unsupported-operation] value_score = (value_score - min_value) / ((max_value - min_value)) return value_score @@ -69,9 +70,9 @@ def qtransform_by_parent_and_siblings( """ chex.assert_shape(node_index, ()) qvalues = tree.qvalues(node_index) - visit_counts = tree.children_visits[node_index] + visit_counts = tree.children_visits[node_index] # pyrefly: ignore[bad-index] chex.assert_rank([qvalues, visit_counts, node_index], [1, 1, 0]) - node_value = tree.node_values[node_index] + node_value = tree.node_values[node_index] # pyrefly: ignore[bad-index] safe_qvalues = jnp.where(visit_counts > 0, qvalues, node_value) chex.assert_equal_shape([safe_qvalues, qvalues]) min_value = jnp.minimum(node_value, jnp.min(safe_qvalues, axis=-1)) @@ -119,12 +120,12 @@ def qtransform_completed_by_mix_value( """ chex.assert_shape(node_index, ()) qvalues = tree.qvalues(node_index) - visit_counts = tree.children_visits[node_index] + visit_counts = tree.children_visits[node_index] # pyrefly: ignore[bad-index] # Computing the mixed value and producing completed_qvalues. - raw_value = tree.raw_values[node_index] + raw_value = tree.raw_values[node_index] # pyrefly: ignore[bad-index] prior_probs = jax.nn.softmax( - tree.children_prior_logits[node_index]) + tree.children_prior_logits[node_index]) # pyrefly: ignore[bad-index] if use_mixed_value: value = _compute_mixed_value( raw_value, diff --git a/mctx/_src/search.py b/mctx/_src/search.py index 18303e8..3415065 100644 --- a/mctx/_src/search.py +++ b/mctx/_src/search.py @@ -98,6 +98,7 @@ def body_fun(sim, loop_state): next_node_index = jnp.where(next_node_index == Tree.UNVISITED, sim + 1, next_node_index) tree = expand( + # pyrefly: ignore[bad-argument-type] params, expand_key, tree, recurrent_fn, parent_index, action, next_node_index) tree = backward(tree, next_node_index) @@ -155,6 +156,7 @@ def body_fun(state): rng_key, action_selection_key = jax.random.split(state.rng_key) action = action_selection_fn(action_selection_key, tree, node_index, state.depth) + # pyrefly: ignore[bad-index] next_node_index = tree.children_index[node_index, action] # The returned action will be visited. depth = state.depth + 1 @@ -164,10 +166,10 @@ def body_fun(state): return _SimulationState( # pytype: disable=wrong-arg-types # jax-types rng_key=rng_key, node_index=node_index, - action=action, - next_node_index=next_node_index, + action=action, # pyrefly: ignore[bad-argument-type] + next_node_index=next_node_index, # pyrefly: ignore[bad-argument-type] depth=depth, - is_continuing=is_continuing) + is_continuing=is_continuing) # pyrefly: ignore[bad-argument-type] node_index = jnp.array(Tree.ROOT_INDEX, dtype=jnp.int32) depth = jnp.zeros((), dtype=tree.children_prior_logits.dtype) @@ -232,7 +234,7 @@ def expand( tree, next_node_index, step.prior_logits, step.value, embedding) # Return updated tree topology. - return tree.replace( + return tree.replace( # pyrefly: ignore[missing-attribute] children_index=batch_update( tree.children_index, next_node_index, parent_index, action), children_rewards=batch_update( @@ -286,6 +288,7 @@ def body_fun(loop_state): return tree, leaf_value, parent leaf_index = jnp.asarray(leaf_index, dtype=jnp.int32) + # pyrefly: ignore[bad-index] loop_state = (tree, tree.node_values[leaf_index], leaf_index) tree, _, _ = jax.lax.while_loop(cond_fun, body_fun, loop_state) @@ -325,6 +328,7 @@ def update_tree_node( chex.assert_shape(prior_logits, (batch_size, tree.num_actions)) # When using max_depth, a leaf can be expanded multiple times. + # pyrefly: ignore[bad-index] new_visit = tree.node_visits[batch_range, node_index] + 1 updates = dict( # pylint: disable=use-dict-literal children_prior_logits=batch_update( @@ -339,7 +343,7 @@ def update_tree_node( lambda t, s: batch_update(t, s, node_index), tree.embeddings, embedding)) - return tree.replace(**updates) + return tree.replace(**updates) # pyrefly: ignore[missing-attribute] def instantiate_tree_from_root( diff --git a/mctx/_src/tests/policies_test.py b/mctx/_src/tests/policies_test.py index 08252a2..549dca9 100644 --- a/mctx/_src/tests/policies_test.py +++ b/mctx/_src/tests/policies_test.py @@ -250,7 +250,7 @@ def test_gumbel_muzero_policy(self): jax.tree.map(lambda x: x[0], policy_output.search_tree), policy_output.search_tree.ROOT_INDEX) self.assertEqual(max_depth, max_found_depth) - self.assertEqual(6, policy_output.search_tree.node_visits[0, leaf]) + self.assertEqual(6, policy_output.search_tree.node_visits[0, leaf]) # pyrefly: ignore[bad-index] def test_gumbel_muzero_policy_without_invalid_actions(self): root_value = jnp.array([-5.0]) diff --git a/mctx/_src/tests/tree_test.py b/mctx/_src/tests/tree_test.py index 7b37659..6b07d4a 100644 --- a/mctx/_src/tests/tree_test.py +++ b/mctx/_src/tests/tree_test.py @@ -106,9 +106,9 @@ def tree_to_pytree(tree: mctx.Tree, batch_i: int = 0): for a_i in range(tree.num_actions): prior = children_prior_probs[batch_i, node_i, a_i] # Index of children, or -1 if not expanded - child_i = int(tree.children_index[batch_i, node_i, a_i]) + child_i = int(tree.children_index[batch_i, node_i, a_i]) # pyrefly: ignore[bad-index] if child_i >= 0: - reward = tree.children_rewards[batch_i, node_i, a_i] + reward = tree.children_rewards[batch_i, node_i, a_i] # pyrefly: ignore[bad-index] child = _create_pynode( tree, batch_i, child_i, prior=prior, action=a_i, reward=reward) nodes[child_i] = child diff --git a/mctx/_src/tree.py b/mctx/_src/tree.py index 43ab99c..e159c82 100644 --- a/mctx/_src/tree.py +++ b/mctx/_src/tree.py @@ -98,11 +98,12 @@ def summary(self) -> SearchSummary: """Extract summary statistics for the root node.""" # Get state and action values for the root nodes. chex.assert_rank(self.node_values, 2) - value = self.node_values[:, Tree.ROOT_INDEX] + value = self.node_values[:, Tree.ROOT_INDEX] # pyrefly: ignore[bad-index] batch_size, = value.shape root_indices = jnp.full((batch_size,), Tree.ROOT_INDEX) qvalues = self.qvalues(root_indices) # Extract visit counts and induced probabilities for the root nodes. + # pyrefly: ignore[bad-index] visit_counts = self.children_visits[:, Tree.ROOT_INDEX].astype(value.dtype) total_counts = jnp.sum(visit_counts, axis=-1, keepdims=True) visit_probs = visit_counts / jnp.maximum(total_counts, 1) @@ -137,6 +138,8 @@ class SearchSummary: def _unbatched_qvalues(tree: Tree, index: int) -> int: chex.assert_rank(tree.children_discounts, 2) return ( # pytype: disable=bad-return-type # numpy-scalars + # pyrefly: ignore[bad-index, bad-return] tree.children_rewards[index] + # pyrefly: ignore[bad-index] + tree.children_discounts[index] * tree.children_values[index] )