lit_sb3.wrappers.mask_velocity
[docs]
module
lit_sb3.wrappers.mask_velocity
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41 | from __future__ import annotations
from typing import ClassVar
import gymnasium as gym
import numpy as np
class MaskVelocityWrapper(gym.ObservationWrapper):
"""Mask velocity terms in observations to make the MDP partially observable.
Adapted from https://github.com/LiuWenlin595/FinalProject.
:param env: Gym environment.
"""
# Supported environments (env_id -> velocity observation indices)
velocity_indices: ClassVar[dict[str, np.ndarray]] = {
"CartPole-v1": np.array([1, 3]),
"MountainCar-v0": np.array([1]),
"MountainCarContinuous-v0": np.array([1]),
"Pendulum-v1": np.array([2]),
"LunarLander-v3": np.array([2, 3, 5]),
"LunarLanderContinuous-v3": np.array([2, 3, 5]),
}
def __init__(self, env: gym.Env) -> None:
super().__init__(env)
assert env.unwrapped.spec is not None
env_id: str = env.unwrapped.spec.id
# By default no masking
self.mask = np.ones_like(env.observation_space.sample())
try:
# Mask velocity
self.mask[self.velocity_indices[env_id]] = 0.0
except KeyError as e:
raise NotImplementedError(f"Velocity masking not implemented for {env_id}") from e
def observation(self, observation: np.ndarray) -> np.ndarray:
return observation * self.mask
|