@@ -27,14 +27,14 @@ SB3 doesn't support `Discrete` and `MultiDiscrete` spaces with `start!=0`. Howev
2727import gymnasium as gym
2828
2929class ShiftWrapper (gym .Wrapper ):
30- """ Allow to use Discrete() action spaces with start!=0"""
31- def __init__ (self , env : gym.Env) -> None :
32- super ().__init__ (env)
33- assert isinstance (env.action_space, gym.spaces.Discrete)
34- self .action_space = gym.spaces.Discrete(env.action_space.n, start = 0 )
35-
36- def step (self , action : int ):
37- return self .env.step(action + self .env.action_space.start)
30+ """ Allow to use Discrete() action spaces with start!=0"""
31+ def __init__ (self , env : gym.Env) -> None :
32+ super ().__init__ (env)
33+ assert isinstance (env.action_space, gym.spaces.Discrete)
34+ self .action_space = gym.spaces.Discrete(env.action_space.n, start = 0 )
35+
36+ def step (self , action : int ):
37+ return self .env.step(action + self .env.action_space.start)
3838```
3939:::
4040
@@ -46,15 +46,15 @@ import numpy as np
4646import gymnasium as gym
4747
4848class ReshapeWrapper (gym .Wrapper ):
49- """ Allow to use MultiDiscrete() action spaces with len(nvec.shape) > 1:"""
50- def __init__ (self , env : gym.Env) -> None :
51- super ().__init__ (env)
52- assert isinstance (env.action_space, gym.spaces.MultiDiscrete)
53- self .original_shape = env.action_space.nvec.shape
54- self .action_space = gym.spaces.MultiDiscrete(env.action_space.nvec.flatten())
55-
56- def step (self , action : np.ndarray):
57- return self .env.step(action.reshape(self .original_shape))
49+ """ Allow to use MultiDiscrete() action spaces with len(nvec.shape) > 1:"""
50+ def __init__ (self , env : gym.Env) -> None :
51+ super ().__init__ (env)
52+ assert isinstance (env.action_space, gym.spaces.MultiDiscrete)
53+ self .original_shape = env.action_space.nvec.shape
54+ self .action_space = gym.spaces.MultiDiscrete(env.action_space.nvec.flatten())
55+
56+ def step (self , action : np.ndarray):
57+ return self .env.step(action.reshape(self .original_shape))
5858```
5959:::
6060
0 commit comments