yapf

google · Dec 20, 2024 · 4174e57 · 4174e57
1 parent 8dd9070
commit 4174e57
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/compiler_opt/rl/imitation_learning/generate_bc_trajectories_lib.py b/compiler_opt/rl/imitation_learning/generate_bc_trajectories_lib.py
@@ -320,8 +320,10 @@ def __init__(
       max_horizon_to_explore=np.inf,
       explore_on_features: Optional[Dict[str, Callable[[tf.Tensor],
                                                        bool]]] = None,
-      obs_action_specs: Optional[Tuple[time_step.TimeStep,
-                                       tensor_spec.BoundedTensorSpec,]] = None,
+      obs_action_specs: Optional[
+        Tuple[time_step.TimeStep,
+        tensor_spec.BoundedTensorSpec,
+      ]] = None,
       reward_key: str = '',
       keep_temps: Optional[str] = None,
       **kwargs,
@@ -751,8 +753,10 @@ def __init__(
       exploration_policy_paths: Optional[str] = None,
       explore_on_features: Optional[Dict[str, Callable[[tf.Tensor],
                                                        bool]]] = None,
-      obs_action_specs: Optional[Tuple[time_step.TimeStep,
-                                       tensor_spec.BoundedTensorSpec,]] = None,
+      obs_action_specs: Optional[Tuple[
+        time_step.TimeStep,
+        tensor_spec.BoundedTensorSpec,
+      ]] = None,
       base_path: Optional[str] = None,
       partitions: List[float] = [
           0.,