fix: fixed type hinting for extra_loss_fns in mappo (#234)

google · Oct 13, 2022 · 2d10e87 · 2d10e87
1 parent 943aa32
commit 2d10e87
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/brax/experimental/composer/training/mappo.py b/brax/experimental/composer/training/mappo.py
@@ -83,7 +83,7 @@ def compute_ppo_loss(
     lambda_: float = 0.95,
     ppo_epsilon: float = 0.3,
     extra_loss_update_ratios: Optional[Dict[str, float]] = None,
-    extra_loss_fns: Optional[Dict[str, Callable[[ppo.StepData],
+    extra_loss_fns: Optional[Dict[str, Callable[[StepData],
                                                 jnp.ndarray]]] = None,
     action_shapes: Dict[str, Dict[str, Any]] = None,
     agent_name: str = None,