Fix lint

HumanCompatibleAI · Oct 7, 2023 · 407b7dd · 407b7dd
1 parent 5bca968
commit 407b7dd
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/src/imitation/scripts/eval_policy.py b/src/imitation/scripts/eval_policy.py
@@ -96,7 +96,7 @@ def eval_policy(
     sample_until = rollout.make_sample_until(eval_n_timesteps, eval_n_episodes)
     post_wrappers = [video_wrapper_factory(log_dir, **video_kwargs)] if videos else None
     with environment.make_venv(  # type: ignore[wrong-keyword-args]
-        post_wrappers=post_wrappers
+        post_wrappers=post_wrappers,
     ) as venv:
         if render:
             venv = InteractiveRender(venv, render_fps)

diff --git a/src/imitation/scripts/ingredients/demonstrations.py b/src/imitation/scripts/ingredients/demonstrations.py
@@ -143,9 +143,9 @@ def _generate_expert_trajs(
         raise ValueError("n_expert_demos must be specified when generating demos.")
 
     logger.info(f"Generating {n_expert_demos} expert trajectories")
-    with environment.make_rollout_venv() as rollout_env:  # type: ignore[wrong-arg-count]
+    with environment.make_rollout_venv() as env:  # type: ignore[wrong-arg-count]
         return rollout.rollout(
-            expert.get_expert_policy(rollout_env),
+            expert.get_expert_policy(env),
             rollout_env,
             rollout.make_sample_until(min_episodes=n_expert_demos),
             rng=_rnd,

diff --git a/src/imitation/scripts/train_rl.py b/src/imitation/scripts/train_rl.py
@@ -100,7 +100,7 @@ def train_rl(
 
     post_wrappers = [lambda env, idx: wrappers.RolloutInfoWrapper(env)]
     with environment.make_venv(  # type: ignore[wrong-keyword-args]
-        post_wrappers=post_wrappers
+        post_wrappers=post_wrappers,
     ) as venv:
         callback_objs = []
         if reward_type is not None: