swmpo.gymnasium_wrapper

  1"""Augment MDPs with state machines using Gymnasium's API."""
  2from collections import defaultdict
  3from gymnasium import Wrapper
  4from swmpo.state_machine import StateMachine
  5from swmpo.state_machine import state_machine_model
  6from swmpo.transition import Transition
  7import numpy as np
  8from gymnasium import spaces
  9from gymnasium import Env
 10import statistics
 11import torch
 12
 13
 14def get_one_hot_encoding(
 15    active_i: int,
 16    total_n: int,
 17) -> list[float]:
 18    """Get the one-hot encoding of the given active class."""
 19    one_hot_vector = [
 20        0.0 if i != active_i else 1.0
 21        for i in range(total_n)
 22    ]
 23    return one_hot_vector
 24
 25
 26def get_augmented_observation(
 27    obs: np.ndarray,
 28    active_i: int,
 29    total_n: int,
 30) -> np.ndarray:
 31    """Augment the observation with a one-hot encoding
 32    vector of the current state of the state machine."""
 33    one_hot_vector = np.array(get_one_hot_encoding(active_i, total_n))
 34    augmented_observation = np.concatenate((
 35        obs,
 36        one_hot_vector,
 37    ))
 38    return augmented_observation
 39
 40
 41def get_exploration_reward(
 42    current_episode_mode_rewards: list[float],
 43    prev_episode_mode_rewards: list[float] | None,
 44    window_size: int,
 45) -> float:
 46    """
 47    Reward the agent for exploring modes where reward has been improving.
 48    - current_reward: reward in current timestep
 49    - mode_rewards: previous rewards in current timestep
 50    """
 51    if len(current_episode_mode_rewards) == 0:
 52        return 0.0
 53    if prev_episode_mode_rewards is None:
 54        return 0.0
 55    if len(prev_episode_mode_rewards) == 0:
 56        return 0.0
 57
 58    # Return the rate of improvement in the window of the last rewards
 59    new_mean = statistics.mean(current_episode_mode_rewards)
 60    old_mean = statistics.mean(prev_episode_mode_rewards)
 61    delta_mean = new_mean - old_mean
 62    return max(0.0, delta_mean)
 63
 64
 65def get_total_reward(
 66    base_reward: float,
 67    current_episode_mode_rewards: list[float],
 68    prev_episode_mode_rewards: list[float] | None,
 69    current_mode: int,
 70    current_episode_visited_modes: set[int],
 71    exploration_window_size: int,
 72    extrinsic_reward_scale: float,
 73) -> float:
 74    # Get exploration reward
 75    exploration_reward = get_exploration_reward(
 76        prev_episode_mode_rewards=prev_episode_mode_rewards,
 77        current_episode_mode_rewards=current_episode_mode_rewards,
 78        window_size=exploration_window_size,
 79    )
 80
 81    # Get extrinsic reward
 82    should_reward = current_mode not in current_episode_visited_modes
 83    extrinsic_reward = extrinsic_reward_scale if should_reward else 0.0
 84
 85    # Get total reward
 86    total_reward = sum((
 87        base_reward,
 88        extrinsic_reward,
 89        exploration_reward
 90    ))
 91    return total_reward
 92
 93
 94def cast_state(state):
 95    if not isinstance(state, torch.Tensor):
 96        if isinstance(state, np.ndarray):
 97            state = torch.from_numpy(state)
 98        else:
 99            state = torch.tensor(state)
100    else:
101        state = state
102    return state
103
104
105class DeepSynthWrapper(Wrapper):
106    """Augment an environment with information from a state machine.
107
108    This wrapper modifies observations and rewards:
109    - reward is augmented providing a constant value for each new
110      node visited during an episode.
111    - observations are augmented with a one-hot encoding of the state
112      of the state machine.
113
114    It is expected the wrapped env returns a dict with a `state` key that
115    contains a `torch.Tensor` that will be used to update the state machine.
116
117    This augmented version of an MDP was proposed in Hasanbeig et al, 2021.
118    """
119
120    def __init__(
121            self,
122            env: Env,
123            state_machine: StateMachine,
124            initial_state_machine_state: int,
125            extrinsic_reward_scale: float,
126            exploration_window_size: int,
127            dt: float,
128            ):
129        super().__init__(env)
130        old_space = env.observation_space
131        assert isinstance(old_space, spaces.Box)
132        encoding_len = len(get_one_hot_encoding(
133            initial_state_machine_state, len(state_machine.local_models),
134        ))
135        low = np.concatenate((
136            old_space.low,
137            np.array([0.0 for _ in range(encoding_len)]),
138        ))
139        high = np.concatenate((
140            old_space.high,
141            np.array([1.0 for _ in range(encoding_len)]),
142        ))
143        self.observation_space = spaces.Box(
144            low=low,
145            high=high,
146            dtype=np.float32,
147        )
148        self.initial_node = initial_state_machine_state
149        self.current_node = initial_state_machine_state
150        self.state_machine = state_machine
151        self.extrinsic_reward_scale = extrinsic_reward_scale
152        self.visited_nodes = set()
153        self.exploration_window_size = exploration_window_size
154        self.dt = dt
155
156        # Keep track of per-mode rewards
157        self.mode_rewards = [defaultdict(list)]
158
159    def reset(self, *args, **kwargs):
160        self.current_node = self.initial_node
161        self.visited_nodes = set()
162        self.mode_rewards.append(defaultdict(list))
163
164        obs, info = self.env.reset(*args, **kwargs)
165        augmented_obs = get_augmented_observation(
166            obs=obs,
167            active_i=self.current_node,
168            total_n=len(self.state_machine.local_models),
169        )
170        info["active_mode"] = self.current_node
171        self.state = info["state"]
172        return augmented_obs, info
173
174    def step(self, action):
175        # Step environment
176        prev_state = self.state
177        obs, reward, terminated, truncated, info = self.env.step(action)
178        info["active_mode"] = self.current_node
179
180        # Step node machine
181        new_state = info["state"]
182        _, new_node = state_machine_model(
183            state_machine=self.state_machine,
184            action=torch.from_numpy(action),
185            state=cast_state(prev_state),
186            current_node=self.current_node,
187            dt=self.dt,
188        )
189
190        # Get total reward
191        if len(self.mode_rewards) < 2:
192            prev_episode_mode_rewards = None
193        else:
194            prev_episode_mode_rewards = self.mode_rewards[-2]
195        total_reward = get_total_reward(
196            base_reward=float(reward),
197            prev_episode_mode_rewards=prev_episode_mode_rewards,
198            current_episode_mode_rewards=self.mode_rewards[-1][new_node],
199            current_mode=new_node,
200            current_episode_visited_modes=self.visited_nodes,
201            exploration_window_size=self.exploration_window_size,
202            extrinsic_reward_scale=self.extrinsic_reward_scale,
203        )
204
205        # Get augmented observation
206        augmented_obs = get_augmented_observation(
207            obs=obs,
208            active_i=new_node,
209            total_n=len(self.state_machine.local_models),
210        )
211
212        # Update per-mode rewards
213        self.mode_rewards[-1][new_node].append(float(reward))
214
215        # Update state machine and visited nodes
216        self.visited_nodes = self.visited_nodes | {new_node}
217        self.current_node = new_node
218
219        self.state = new_state
220
221        return augmented_obs, total_reward, terminated, truncated, info