From fb57a6b80c1936b5f9fc40e2ce2ea2c075d325b1 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Thu, 23 Jan 2020 11:20:12 +0100 Subject: [PATCH] Update docstring --- torchy_baselines/common/vec_env/vec_frame_stack.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/torchy_baselines/common/vec_env/vec_frame_stack.py b/torchy_baselines/common/vec_env/vec_frame_stack.py index 562c525..6676162 100644 --- a/torchy_baselines/common/vec_env/vec_frame_stack.py +++ b/torchy_baselines/common/vec_env/vec_frame_stack.py @@ -3,18 +3,18 @@ import warnings import numpy as np from gym import spaces -from torchy_baselines.common.vec_env.base_vec_env import VecEnvWrapper +from torchy_baselines.common.vec_env.base_vec_env import VecEnv, VecEnvWrapper class VecFrameStack(VecEnvWrapper): """ Frame stacking wrapper for vectorized environment - :param venv: (VecEnv) the vectorized environment to wrap - :param n_stack: (int) Number of frames to stack + :param venv: the vectorized environment to wrap + :param n_stack: Number of frames to stack """ - def __init__(self, venv, n_stack): + def __init__(self, venv: VecEnv, n_stack: int): self.venv = venv self.n_stack = n_stack wrapped_obs_space = venv.observation_space