From 21b3d9029b38c67b15580a28f93dd0caa5036f90 Mon Sep 17 00:00:00 2001 From: Dqn Zoo Team Date: Tue, 21 Feb 2023 23:13:50 +0000 Subject: [PATCH] Fix or ignore some pytype errors related to jnp.ndarray == jax.Array. PiperOrigin-RevId: 511317999 --- dqn_zoo/iqn/agent.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/dqn_zoo/iqn/agent.py b/dqn_zoo/iqn/agent.py index e1dc997..80ea8fb 100644 --- a/dqn_zoo/iqn/agent.py +++ b/dqn_zoo/iqn/agent.py @@ -76,7 +76,9 @@ def select_action(rng_key, network_params, s_t): tau_t = _sample_tau(tau_key, (1, tau_samples)) q_t = network.apply( network_params, apply_key, IqnInputs(s_t[None, ...], tau_t) - ).q_values[0] + ).q_values[ + 0 + ] # pytype: disable=wrong-arg-types # jax-ndarray a_t = distrax.EpsilonGreedy(q_t, exploration_epsilon).sample( seed=policy_key ) @@ -188,15 +190,15 @@ def loss_fn(online_params, target_params, transitions, rng_key): _, *apply_keys = jax.random.split(rng_key, 4) dist_q_tm1 = network.apply( online_params, apply_keys[0], IqnInputs(transitions.s_tm1, tau_tm1) - ).q_dist + ).q_dist # pytype: disable=wrong-arg-types # jax-ndarray dist_q_t_selector = network.apply( target_params, apply_keys[1], IqnInputs(transitions.s_t, tau_t_selector), - ).q_dist + ).q_dist # pytype: disable=wrong-arg-types # jax-ndarray dist_q_target_t = network.apply( target_params, apply_keys[2], IqnInputs(transitions.s_t, tau_t) - ).q_dist + ).q_dist # pytype: disable=wrong-arg-types # jax-ndarray losses = _batch_quantile_q_learning( dist_q_tm1, tau_tm1, @@ -229,7 +231,9 @@ def select_action(rng_key, network_params, s_t, exploration_epsilon): tau_t = _sample_tau(sample_key, (1, tau_samples_policy)) q_t = network.apply( network_params, apply_key, IqnInputs(s_t[None, ...], tau_t) - ).q_values[0] + ).q_values[ + 0 + ] # pytype: disable=wrong-arg-types # jax-ndarray a_t = distrax.EpsilonGreedy(q_t, exploration_epsilon).sample( seed=policy_key )