Closed capnspacehook closed 11 months ago
Hello, thanks for the bug report. I guess the issue comes from a flatten layer which is not applied in SBX.
A quick fix is to use a VecFrameStack
instead (it stacks on the last axis instead of the first):
from stable_baselines3.common.vec_env import DummyVecEnv, VecFrameStack
vec_env = DummyVecEnv([lambda: env])
vec_env = VecFrameStack(vec_env, 4)
To reproduce with a even more minimal code:
from typing import Any, List, Tuple
import gymnasium as gym
from gymnasium.spaces import Box, Discrete
from sbx import PPO
class MyEnv(gym.Env):
def __init__(self) -> None:
self.observation_space = Box(0, 1, shape=(2, 1), dtype="float32")
self.action_space = Discrete(15)
def step(self, action: Any) -> Tuple[Any, float, bool, bool, dict]:
return self.observation_space.sample(), 0.0, False, False, {}
def reset(
self, *, seed: int | None = None, options: dict | None = None
) -> Tuple[Any, dict]:
super().reset(seed=seed, options=options)
return self.observation_space.sample(), {}
def render(self) -> Any | List[Any] | None:
pass
PPO("MlpPolicy", MyEnv()).learn(total_timesteps=1000)
I've pushed a fix in https://github.com/araffin/sbx/pull/18, you should be able to upgrade to sbx 0.9.0 soon =)
🤖 Custom Gym Environment
Describe the bug
When using
gymnasium.wrappers.frame_stack.FrameStack
with a simple custom env, I get an exception when an action is being chosen instep
.Code example
 System Info
sbx
at the latest commit was installed usingpip
:pip install git+https://github.com/araffin/sbx
 Checklist