simonbb 0.0.3__tar.gz → 0.0.4__tar.gz

{simonbb-0.0.3 → simonbb-0.0.4}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: simonbb
-Version: 0.0.3
-Summary: A small package to facilitate teaching AI/ML courses
+Version: 0.0.4
+Summary: A small package to facilitate teaching AI courses
 Project-URL: Homepage, https://simonbb.com
 Author-email: Huaxia Rui <huaxia.rui@simon.rochester.edu>
 License-Expression: MIT
@@ -9,6 +9,7 @@ License-File: LICENSE
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.9
+Requires-Dist: gymnasium~=1.3.0
 Requires-Dist: langchain-chroma~=1.1.0
 Requires-Dist: langchain-community~=0.4.1
 Requires-Dist: langchain-core~=1.2.23
@@ -17,6 +18,8 @@ Requires-Dist: langchain~=1.2.13
 Requires-Dist: langgraph~=1.1.3
 Requires-Dist: matplotlib~=3.10.8
 Requires-Dist: numpy~=2.4.0
+Requires-Dist: pytorch-ignite~=0.5.4
+Requires-Dist: pyyaml~=6.0.3
 Requires-Dist: sentence-transformers~=5.3.0
 Requires-Dist: torchvision~=0.24.1
 Requires-Dist: torch~=2.9.1

{simonbb-0.0.3 → simonbb-0.0.4}/pyproject.toml

@@ -4,11 +4,11 @@ build-backend = "hatchling.build"
 
 [project]
 name = "simonbb"
-version = "0.0.3"
+version = "0.0.4"
 authors = [
 	{ name="Huaxia Rui", email="huaxia.rui@simon.rochester.edu" },
 ]
-description = "A small package to facilitate teaching AI/ML courses"
+description = "A small package to facilitate teaching AI courses"
 readme = "README.md"
 requires-python = ">=3.9"
 classifiers = [
@@ -30,6 +30,9 @@ dependencies = [
 	"langchain-chroma~=1.1.0",
 	"langchain-huggingface~=1.2.1",
 	"sentence-transformers~=5.3.0",
+	"gymnasium~=1.3.0",
+	"pytorch-ignite~=0.5.4",
+	"pyyaml~=6.0.3"
 ]
 [project.urls]
 Homepage = "https://simonbb.com"

simonbb-0.0.4/src/simonbb/cis433/actions.py

@@ -0,0 +1,76 @@
+import abc
+import numpy as np
+import typing as tt
+
+class ActionSelector(abc.ABC):
+    """
+    Abstract class which converts scores to the actions
+    """
+    @abc.abstractmethod
+    def __call__(self, scores: np.ndarray) -> np.ndarray:
+        ...
+
+
+class ArgmaxActionSelector(ActionSelector):
+    """
+    Selects actions using argmax
+    """
+    def __call__(self, scores: np.ndarray) -> np.ndarray:
+        return np.argmax(scores, axis=1)
+
+
+class EpsilonGreedyActionSelector(ActionSelector):
+    def __init__(self, epsilon: float = 0.05,
+                 selector: tt.Optional[ActionSelector] = None):
+        self._epsilon = epsilon
+        self.selector = selector if selector is not None \
+            else ArgmaxActionSelector()
+
+    @property
+    def epsilon(self) -> float:
+        return self._epsilon
+
+    @epsilon.setter
+    def epsilon(self, value: float):
+        if value < 0.0 or value > 1.0:
+            raise ValueError("Epsilon has to be between 0 and 1")
+        self._epsilon = value
+
+    def __call__(self, scores: np.ndarray) -> np.ndarray:
+        assert len(scores.shape) == 2
+        batch_size, n_actions = scores.shape
+        actions = self.selector(scores)
+        mask = np.random.random(size=batch_size) < self.epsilon
+        rand_actions = np.random.choice(n_actions, sum(mask))
+        actions[mask] = rand_actions
+        return actions
+
+
+class ProbabilityActionSelector(ActionSelector):
+    """
+    Converts probabilities of actions into action by sampling them
+    """
+    def __call__(self, probs: np.ndarray) -> np.ndarray:
+        actions = []
+        for prob in probs:
+            actions.append(np.random.choice(len(prob), p=prob))
+        return np.array(actions)
+
+
+class EpsilonTracker:
+    """
+    Updates epsilon according to linear schedule
+    """
+    def __init__(self, selector: EpsilonGreedyActionSelector,
+                 eps_start: tt.Union[int, float],
+                 eps_final: tt.Union[int, float],
+                 eps_frames: int):
+        self.selector = selector
+        self.eps_start = eps_start
+        self.eps_final = eps_final
+        self.eps_frames = eps_frames
+        self.frame(0)
+
+    def frame(self, frame: int):
+        eps = self.eps_start - frame / self.eps_frames
+        self.selector.epsilon = max(self.eps_final, eps)

simonbb-0.0.4/src/simonbb/cis433/agent.py

@@ -0,0 +1,193 @@
+"""
+Agent is something which converts states into actions and has state
+"""
+import abc
+import copy
+import numpy as np
+import torch
+from torch import nn
+import torch.nn.functional as F
+import typing as tt
+
+from . import actions
+
+States = tt.List[np.ndarray] | np.ndarray
+AgentStates = tt.List[tt.Any]
+Preprocessor = tt.Callable[[States], torch.Tensor]
+
+CPU_DEVICE = torch.device("cpu")
+
+
+def default_states_preprocessor(states: States) -> torch.Tensor:
+    """
+    Convert list of states into the form suitable for model
+    :param states: list of numpy arrays with states or numpy array
+    :return: torch.Tensor
+    """
+    if isinstance(states, list):
+        if len(states) == 1:
+            np_states = np.expand_dims(states[0], 0)
+        else:
+            np_states = np.asarray([np.asarray(s) for s in states])
+    else:
+        np_states = states
+    return torch.as_tensor(np_states)
+
+
+def float32_preprocessor(states: States):
+    np_states = np.array(states, dtype=np.float32)
+    return torch.as_tensor(np_states)
+
+
+class BaseAgent(abc.ABC):
+    """
+    Base Agent, sharing most of logic with concrete agent implementations.
+    """
+    def initial_state(self) -> tt.Optional[tt.Any]:
+        """
+        Should create initial empty state for the agent. It will be
+        called for the start of the episode
+        :return: Anything agent want to remember
+        """
+        return None
+
+    @abc.abstractmethod
+    def __call__(self, states: States, agent_states: AgentStates) -> tt.Tuple[np.ndarray, AgentStates]:
+        ...
+
+
+class NNAgent(BaseAgent):
+    """
+    Network-based agent
+    """
+    def __init__(self, model: nn.Module, action_selector: actions.ActionSelector,
+                 device: torch.device, preprocessor: Preprocessor):
+        """
+        Constructor of base agent
+        :param model: model to be used
+        :param action_selector: action selector
+        :param device: device for tensors
+        :param preprocessor: states preprocessor
+        """
+        self.model = model
+        self.action_selector = action_selector
+        self.device = device
+        self.preprocessor = preprocessor
+
+    @abc.abstractmethod
+    def _net_filter(self, net_out: tt.Any, agent_states: AgentStates) -> \
+            tt.Tuple[torch.Tensor, AgentStates]:
+        """
+        Internal method, processing network output and states into selector's input and new states
+        :param net_out: output from the network
+        :param agent_states: agent states
+        :return: tuple with tensor to be fed into selector and new states
+        """
+        ...
+
+    @torch.no_grad()
+    def __call__(self, states: States, agent_states: AgentStates = None) -> tt.Tuple[np.ndarray, AgentStates]:
+        """
+        Convert observations and states into actions to take
+        :param states: list of environment states to process
+        :param agent_states: list of states with the same length as observations
+        :return: tuple of actions, states
+        """
+        if agent_states is None:
+            agent_states = [None] * len(states)
+        if self.preprocessor is not None:
+            states = self.preprocessor(states)
+            if torch.is_tensor(states):
+                states = states.to(self.device)
+        q_v = self.model(states)
+        q_v, new_states = self._net_filter(q_v, agent_states)
+        q = q_v.data.cpu().numpy()
+        actions = self.action_selector(q)
+        return actions, new_states
+
+
+class DQNAgent(NNAgent):
+    """
+    DQNAgent is a memoryless DQN agent which calculates Q values
+    from the observations and  converts them into the actions using action_selector
+    """
+    def __init__(self, model: nn.Module, action_selector: actions.ActionSelector,
+                 device: torch.device = CPU_DEVICE,
+                 preprocessor: Preprocessor = default_states_preprocessor):
+        super().__init__(model, action_selector=action_selector, device=device,
+                         preprocessor=preprocessor)
+
+    # not needed in DQN - we don't process Q-values returned
+    def _net_filter(self, net_out: tt.Any, agent_states: AgentStates) -> \
+            tt.Tuple[torch.Tensor, AgentStates]:
+        assert torch.is_tensor(net_out)
+        return net_out, agent_states
+
+
+class TargetNet:
+    """
+    Wrapper around model which provides copy of it instead of trained weights
+    """
+    def __init__(self, model: nn.Module):
+        self.model = model
+        self.target_model = copy.deepcopy(model)
+
+    def sync(self):
+        self.target_model.load_state_dict(self.model.state_dict())
+
+    def alpha_sync(self, alpha):
+        """
+        Blend params of target net with params from the model
+        :param alpha:
+        """
+        assert isinstance(alpha, float)
+        assert 0.0 < alpha <= 1.0
+        state = self.model.state_dict()
+        tgt_state = self.target_model.state_dict()
+        for k, v in state.items():
+            tgt_state[k] = tgt_state[k] * alpha + (1 - alpha) * v
+        self.target_model.load_state_dict(tgt_state)
+
+
+class PolicyAgent(NNAgent):
+    """
+    Policy agent gets action probabilities from the model and samples actions from it
+    """
+    def __init__(self, model: nn.Module,
+                 action_selector: actions.ActionSelector = actions.ProbabilityActionSelector(),
+                 device: torch.device = CPU_DEVICE, apply_softmax: bool = False,
+                 preprocessor: Preprocessor = default_states_preprocessor):
+        super().__init__(model=model, action_selector=action_selector, device=device,
+                         preprocessor=preprocessor)
+        self.apply_softmax = apply_softmax
+
+    def _net_filter(self, net_out: tt.Any, agent_states: AgentStates) -> \
+            tt.Tuple[torch.Tensor, AgentStates]:
+        assert torch.is_tensor(net_out)
+        if self.apply_softmax:
+            return F.softmax(net_out, dim=1), agent_states
+        return net_out, agent_states
+
+
+class ActorCriticAgent(NNAgent):
+    """
+    Policy agent which returns policy and value tensors from observations. Value are stored in agent's state
+    and could be reused for rollouts calculations by ExperienceSource.
+    """
+    def __init__(self, model: nn.Module,
+                 action_selector: actions.ActionSelector = actions.ProbabilityActionSelector(),
+                 device: torch.device = CPU_DEVICE, apply_softmax: bool = False,
+                 preprocessor: Preprocessor = default_states_preprocessor):
+        super().__init__(model=model, action_selector=action_selector, device=device, preprocessor=preprocessor)
+        self.apply_softmax = apply_softmax
+
+    def _net_filter(self, net_out: tt.Any, agent_states: AgentStates) -> \
+            tt.Tuple[torch.Tensor, AgentStates]:
+        assert isinstance(net_out, tuple)
+        policy_v, values_v = net_out
+        assert torch.is_tensor(policy_v)
+        assert torch.is_tensor(values_v)
+        agent_states = values_v.data.squeeze().cpu().numpy().tolist()
+        if self.apply_softmax:
+            return F.softmax(policy_v, dim=1), agent_states
+        return policy_v, agent_states

simonbb-0.0.4/src/simonbb/cis433/experience.py

@@ -0,0 +1,604 @@
+import gymnasium as gym
+import torch
+import random
+import collections
+import typing as tt
+from dataclasses import dataclass
+
+import numpy as np
+
+from collections import deque
+
+from .agent import BaseAgent
+
+State = np.ndarray
+Action = int
+
+
+@dataclass(frozen=True)
+class Experience:
+    state: State
+    action: Action
+    reward: float
+    done_trunc: bool
+
+
+class ExperienceSource:
+    """
+    Simple n-step experience source using single or multiple environments
+
+    Every experience contains n list of Experience entries
+    """
+    Item = tt.Tuple[Experience, ...]
+
+    def __init__(self, env: gym.Env | tt.Collection[gym.Env], agent: BaseAgent,
+                 steps_count: int = 2, steps_delta: int = 1,
+                 env_seed: tt.Optional[int] = None):
+        """
+        Create simple experience source
+        :param env: environment or list of environments to be used
+        :param agent: callable to convert batch of states into actions to take
+        :param steps_count: count of steps to track for every experience chain
+        :param steps_delta: how many steps to do between experience items
+        :param env_seed: seed to be used in Env.reset() call
+        """
+        assert steps_count >= 1
+        if isinstance(env, (list, tuple)):
+            self.pool = env
+            # do the check for the multiple copies passed
+            ids = set(id(e) for e in env)
+            if len(ids) < len(env):
+                raise ValueError("You passed single environment instance multiple times")
+        else:
+            self.pool = [env]
+        self.agent = agent
+        self.steps_count = steps_count
+        self.steps_delta = steps_delta
+        self.total_rewards = []
+        self.total_steps = []
+        self.agent_states = [agent.initial_state() for _ in self.pool]
+        self.env_seed = env_seed
+
+    def __iter__(self) -> tt.Generator[Item, None, None]:
+        states, histories, cur_rewards, cur_steps = [], [], [], []
+        for env in self.pool:
+            if self.env_seed is not None:
+                obs, _ = env.reset(seed=self.env_seed)
+            else:
+                obs, _ = env.reset()
+            states.append(obs)
+            histories.append(deque(maxlen=self.steps_count))
+            cur_rewards.append(0.0)
+            cur_steps.append(0)
+
+        iter_idx = 0
+        while True:
+            actions, self.agent_states = self.agent(states, self.agent_states)
+            for idx, env in enumerate(self.pool):
+                state = states[idx]
+                action = actions[idx]
+                history = histories[idx]
+                next_state, r, is_done, is_tr, _ = env.step(action)
+                cur_rewards[idx] += r
+                cur_steps[idx] += 1
+                history.append(Experience(state=state, action=action, reward=r, done_trunc=is_done or is_tr))
+                if len(history) == self.steps_count and iter_idx % self.steps_delta == 0:
+                    yield tuple(history)
+                states[idx] = next_state
+                if is_done or is_tr:
+                    # generate tail of history
+                    if 0 < len(history) < self.steps_count:
+                        yield tuple(history)
+                    while len(history) > 1:
+                        history.popleft()
+                        yield tuple(history)
+                    self.total_rewards.append(cur_rewards[idx])
+                    self.total_steps.append(cur_steps[idx])
+                    cur_rewards[idx] = 0.0
+                    cur_steps[idx] = 0
+                    if self.env_seed is not None:
+                        states[idx], _ = env.reset(seed=self.env_seed)
+                    else:
+                        states[idx], _ = env.reset()
+                    self.agent_states[idx] = self.agent.initial_state()
+                    history.clear()
+            iter_idx += 1
+
+    def pop_total_rewards(self) -> tt.List[float]:
+        r = self.total_rewards
+        if r:
+            self.total_rewards = []
+            self.total_steps = []
+        return r
+
+    def pop_rewards_steps(self) -> tt.List[tt.Tuple[float, int]]:
+        res = list(zip(self.total_rewards, self.total_steps))
+        if res:
+            self.total_rewards, self.total_steps = [], []
+        return res
+
+
+@dataclass(frozen=True)
+class ExperienceFirstLast:
+    state: State
+    action: Action
+    reward: float
+    last_state: tt.Optional[State]
+
+
+class ExperienceSourceFirstLast(ExperienceSource):
+    """
+    This is a wrapper around ExperienceSource to prevent storing full trajectory in replay buffer when we need
+    only first and last states. For every trajectory piece it calculates discounted reward and emits only first
+    and last states and action taken in the first state.
+
+    If we have partial trajectory at the end of episode, last_state will be None
+    """
+    def __init__(self, env: gym.Env, agent: BaseAgent, gamma: float,
+                 steps_count: int = 1, steps_delta: int = 1, env_seed: tt.Optional[int] = None):
+        super(ExperienceSourceFirstLast, self).__init__(env, agent, steps_count+1, steps_delta, env_seed=env_seed)
+        self.gamma = gamma
+        self.steps = steps_count
+
+    def __iter__(self) -> tt.Generator[ExperienceFirstLast, None, None]:
+        for exp in super(ExperienceSourceFirstLast, self).__iter__():
+            if exp[-1].done_trunc and len(exp) <= self.steps:
+                last_state = None
+                elems = exp
+            else:
+                last_state = exp[-1].state
+                elems = exp[:-1]
+            total_reward = 0.0
+            for e in reversed(elems):
+                total_reward *= self.gamma
+                total_reward += e.reward
+            yield ExperienceFirstLast(state=exp[0].state, action=exp[0].action,
+                                      reward=total_reward, last_state=last_state)
+
+
+def vector_rewards(rewards: tt.Deque[np.ndarray], dones: tt.Deque[np.ndarray], gamma: float) -> np.ndarray:
+    """
+    Calculate rewards from vectorized environment for given amount of steps.
+    :param rewards: deque with observed rewards
+    :param dones: deque with bool flags indicating end of episode
+    :param gamma: gamma constant
+    :return: vector with accumulated rewards
+    """
+    res = np.zeros(rewards[0].shape[0], dtype=np.float32)
+    for r, d in zip(reversed(rewards), reversed(dones)):
+        res *= gamma * (1. - d)
+        res += r
+    return res
+
+
+class VectorExperienceSourceFirstLast(ExperienceSource):
+    """
+    ExperienceSourceFirstLast which supports VectorEnv from Gymnasium.
+    """
+    def __init__(self, env: gym.vector.VectorEnv, agent: BaseAgent,
+                 gamma: float, steps_count: int = 1, env_seed: tt.Optional[int] = None,
+                 unnest_data: bool = True):
+        """
+        Construct vectorized version of ExperienceSourceFirstLast
+        :param env: vectorized environment
+        :param agent: agent to use
+        :param gamma: gamma for reward calculation
+        :param steps_count: count of steps
+        :param env_seed: seed for environments reset
+        :param unnest_data: should we unnest data in the iterator. If True (default)
+        ExperienceFirstLast will be yielded sequentially. If False, we'll keep them in a list as we
+        got them from env vector.
+        """
+        super().__init__(env, agent, steps_count+1, steps_delta=1, env_seed=env_seed)
+        self.env = env
+        self.gamma = gamma
+        self.steps = steps_count
+        self.unnest_data = unnest_data
+        self.agent_state = self.agent_states[0]
+
+    def _iter_env_idx_obs_next(self, b_obs, b_next_obs) -> tt.Generator[tt.Tuple[
+        int, tt.Any, tt.Any
+    ], None, None]:
+        """
+        Iterate over individual environment observations and next observations.
+        Take into account Tuple observation space (which is handled specially in Vectorized envs)
+        :param b_obs: vectorized observations
+        :param b_next_obs: vectorized next observations
+        :yields: Tuple of index, observation and next observation
+        """
+        if isinstance(self.env.single_observation_space, gym.spaces.Tuple):
+            obs_iter = zip(*b_obs)
+            next_obs_iter = zip(*b_next_obs)
+        else:
+            obs_iter = b_obs
+            next_obs_iter = b_next_obs
+        yield from zip(range(self.env.num_envs), obs_iter, next_obs_iter)
+
+    def __iter__(self) -> tt.Generator[tt.List[ExperienceFirstLast] | ExperienceFirstLast, None, None]:
+        q_states = collections.deque(maxlen=self.steps+1)
+        q_actions = collections.deque(maxlen=self.steps+1)
+        q_rewards = collections.deque(maxlen=self.steps+1)
+        q_dones = collections.deque(maxlen=self.steps+1)
+        total_rewards = np.zeros(self.env.num_envs, dtype=np.float32)
+        total_steps = np.zeros_like(total_rewards, dtype=np.int64)
+
+        if self.env_seed is not None:
+            obs, _ = self.env.reset(seed=self.env_seed)
+        else:
+            obs, _ = self.env.reset()
+        env_indices = np.arange(self.env.num_envs)
+
+        while True:
+            q_states.append(obs)
+            actions, self.agent_state = self.agent(obs, self.agent_state)
+            q_actions.append(actions)
+            next_obs, r, is_done, is_tr, _ = self.env.step(actions)
+            total_rewards += r
+            total_steps += 1
+            done_or_tr = is_done | is_tr
+            q_rewards.append(r)
+            q_dones.append(done_or_tr)
+
+            # process environments which are done at this step
+            if done_or_tr.any():
+                indices = env_indices[done_or_tr]
+                self.total_steps.extend(total_steps[indices])
+                self.total_rewards.extend(total_rewards[indices])
+                total_steps[indices] = 0
+                total_rewards[indices] = 0.0
+
+            if len(q_states) == q_states.maxlen:
+                # enough data for calculation
+                results = []
+                rewards = vector_rewards(q_rewards, q_dones, self.gamma)
+                for i, e_obs, e_next_obs in self._iter_env_idx_obs_next(q_states[0], next_obs):
+                    # if anywhere in the trajectory we have ended episode flag,
+                    # the last state will be None
+                    ep_ended = any(map(lambda d: d[i], q_dones))
+                    last_state = e_next_obs if not ep_ended else None
+                    results.append(ExperienceFirstLast(
+                        state=e_obs,
+                        action=q_actions[0][i],
+                        reward=rewards[i],
+                        last_state=last_state,
+                    ))
+                if self.unnest_data:
+                    yield from results
+                else:
+                    yield results
+            obs = next_obs
+
+
+def discount_with_dones(rewards, dones, gamma):
+    discounted = []
+    r = 0
+    for reward, done in zip(rewards[::-1], dones[::-1]):
+        r = reward + gamma*r*(1.-done)
+        discounted.append(r)
+    return discounted[::-1]
+
+
+class ExperienceSourceRollouts:
+    """
+    TODO: need to be updated
+    N-step rollout experience source following A3C rollouts scheme. Have to be used with agent,
+    keeping the value in its state (for example, agent.ActorCriticAgent).
+
+    Yields batches of num_envs * n_steps samples with the following arrays:
+    1. observations
+    2. actions
+    3. discounted rewards, with values approximation
+    4. values
+    """
+    def __init__(self, env, agent, gamma, steps_count=5):
+        """
+        Constructs the rollout experience source
+        :param env: environment or list of environments to be used
+        :param agent: callable to convert batch of states into actions
+        :param steps_count: how many steps to perform rollouts
+        """
+        assert isinstance(env, (gym.Env, list, tuple))
+        assert isinstance(agent, BaseAgent)
+        assert isinstance(gamma, float)
+        assert isinstance(steps_count, int)
+        assert steps_count >= 1
+
+        if isinstance(env, (list, tuple)):
+            self.pool = env
+        else:
+            self.pool = [env]
+        self.agent = agent
+        self.gamma = gamma
+        self.steps_count = steps_count
+        self.total_rewards = []
+        self.total_steps = []
+
+    def __iter__(self):
+        pool_size = len(self.pool)
+        states = [np.array(e.reset()) for e in self.pool]
+        mb_states = np.zeros((pool_size, self.steps_count) + states[0].shape, dtype=states[0].dtype)
+        mb_rewards = np.zeros((pool_size, self.steps_count), dtype=np.float32)
+        mb_values = np.zeros((pool_size, self.steps_count), dtype=np.float32)
+        mb_actions = np.zeros((pool_size, self.steps_count), dtype=np.int64)
+        mb_dones = np.zeros((pool_size, self.steps_count), dtype=np.bool)
+        total_rewards = [0.0] * pool_size
+        total_steps = [0] * pool_size
+        agent_states = None
+        step_idx = 0
+
+        while True:
+            actions, agent_states = self.agent(states, agent_states)
+            rewards = []
+            dones = []
+            new_states = []
+            for env_idx, (e, action) in enumerate(zip(self.pool, actions)):
+                o, r, done, _ = e.step(action)
+                total_rewards[env_idx] += r
+                total_steps[env_idx] += 1
+                if done:
+                    o = e.reset()
+                    self.total_rewards.append(total_rewards[env_idx])
+                    self.total_steps.append(total_steps[env_idx])
+                    total_rewards[env_idx] = 0.0
+                    total_steps[env_idx] = 0
+                new_states.append(np.array(o))
+                dones.append(done)
+                rewards.append(r)
+            # we need an extra step to get values approximation for rollouts
+            if step_idx == self.steps_count:
+                # calculate rollout rewards
+                for env_idx, (env_rewards, env_dones, last_value) in enumerate(zip(mb_rewards, mb_dones, agent_states)):
+                    env_rewards = env_rewards.tolist()
+                    env_dones = env_dones.tolist()
+                    if not env_dones[-1]:
+                        env_rewards = discount_with_dones(env_rewards + [last_value], env_dones + [False], self.gamma)[:-1]
+                    else:
+                        env_rewards = discount_with_dones(env_rewards, env_dones, self.gamma)
+                    mb_rewards[env_idx] = env_rewards
+                yield mb_states.reshape((-1,) + mb_states.shape[2:]), mb_rewards.flatten(), mb_actions.flatten(), mb_values.flatten()
+                step_idx = 0
+            mb_states[:, step_idx] = states
+            mb_rewards[:, step_idx] = rewards
+            mb_values[:, step_idx] = agent_states
+            mb_actions[:, step_idx] = actions
+            mb_dones[:, step_idx] = dones
+            step_idx += 1
+            states = new_states
+
+    def pop_total_rewards(self):
+        r = self.total_rewards
+        if r:
+            self.total_rewards = []
+            self.total_steps = []
+        return r
+
+    def pop_rewards_steps(self):
+        res = list(zip(self.total_rewards, self.total_steps))
+        if res:
+            self.total_rewards, self.total_steps = [], []
+        return res
+
+
+class ExperienceSourceBuffer:
+    """
+    The same as ExperienceSource, but takes episodes from the buffer
+    """
+    def __init__(self, buffer, steps_count=1):
+        """
+        Create buffered experience source
+        :param buffer: list of episodes, each is a list of Experience object
+        :param steps_count: count of steps in every entry
+        """
+        self.update_buffer(buffer)
+        self.steps_count = steps_count
+
+    def update_buffer(self, buffer):
+        self.buffer = buffer
+        self.lens = list(map(len, buffer))
+
+    def __iter__(self):
+        """
+        Infinitely sample episode from the buffer and then sample item offset
+        """
+        while True:
+            episode = random.randrange(len(self.buffer))
+            ofs = random.randrange(self.lens[episode] - self.steps_count - 1)
+            yield self.buffer[episode][ofs:ofs+self.steps_count]
+
+
+class ExperienceReplayBuffer:
+    def __init__(self, experience_source: tt.Optional[ExperienceSource], buffer_size: int):
+        self.experience_source_iter = None if experience_source is None else iter(experience_source)
+        self.buffer: tt.List[ExperienceSource.Item] = []
+        self.capacity = buffer_size
+        self.pos = 0
+
+    def __len__(self):
+        return len(self.buffer)
+
+    def __iter__(self) -> tt.Iterator[ExperienceSource.Item]:
+        return iter(self.buffer)
+
+    def sample(self, batch_size: int) -> tt.List[ExperienceSource.Item]:
+        """
+        Get one random batch from experience replay
+        :param batch_size: size of the batch to sample
+        :return: list of experience entries
+        """
+        if len(self.buffer) <= batch_size:
+            return self.buffer
+        # Warning: replace=False makes random.choice O(n)
+        keys = np.random.choice(len(self.buffer), batch_size, replace=True)
+        return [self.buffer[key] for key in keys]
+
+    def _add(self, sample: ExperienceSource.Item):
+        if len(self.buffer) < self.capacity:
+            self.buffer.append(sample)
+        else:
+            self.buffer[self.pos] = sample
+        self.pos = (self.pos + 1) % self.capacity
+
+    def populate(self, samples: int):
+        """
+        Populates samples into the buffer
+        :param samples: how many samples to populate
+        """
+        for _ in range(samples):
+            entry = next(self.experience_source_iter)
+            self._add(entry)
+
+
+class PrioReplayBufferNaive:
+    def __init__(self, exp_source, buf_size, prob_alpha=0.6):
+        self.exp_source_iter = iter(exp_source)
+        self.prob_alpha = prob_alpha
+        self.capacity = buf_size
+        self.pos = 0
+        self.buffer = []
+        self.priorities = np.zeros((buf_size, ), dtype=np.float32)
+
+    def __len__(self):
+        return len(self.buffer)
+
+    def populate(self, count):
+        max_prio = self.priorities.max() if self.buffer else 1.0
+        for _ in range(count):
+            sample = next(self.exp_source_iter)
+            if len(self.buffer) < self.capacity:
+                self.buffer.append(sample)
+            else:
+                self.buffer[self.pos] = sample
+            self.priorities[self.pos] = max_prio
+            self.pos = (self.pos + 1) % self.capacity
+
+    def sample(self, batch_size, beta=0.4):
+        if len(self.buffer) == self.capacity:
+            prios = self.priorities
+        else:
+            prios = self.priorities[:self.pos]
+        probs = np.array(prios, dtype=np.float32) ** self.prob_alpha
+
+        probs /= probs.sum()
+        indices = np.random.choice(len(self.buffer), batch_size, p=probs, replace=True)
+        samples = [self.buffer[idx] for idx in indices]
+        total = len(self.buffer)
+        weights = (total * probs[indices]) ** (-beta)
+        weights /= weights.max()
+        return samples, indices, np.array(weights, dtype=np.float32)
+
+    def update_priorities(self, batch_indices, batch_priorities):
+        for idx, prio in zip(batch_indices, batch_priorities):
+            self.priorities[idx] = prio
+
+
+
+class BatchPreprocessor:
+    """
+    Abstract preprocessor class descendants to which converts experience
+    batch to form suitable to learning.
+    """
+    def preprocess(self, batch):
+        raise NotImplementedError
+
+
+class QLearningPreprocessor(BatchPreprocessor):
+    """
+    Supports SimpleDQN, TargetDQN, DoubleDQN and can additionally feed TD-error back to
+    experience replay buffer.
+
+    To use different modes, use appropriate class method
+    """
+    def __init__(self, model, target_model, use_double_dqn=False, batch_td_error_hook=None, gamma=0.99, device="cpu"):
+        self.model = model
+        self.target_model = target_model
+        self.use_double_dqn = use_double_dqn
+        self.batch_dt_error_hook = batch_td_error_hook
+        self.gamma = gamma
+        self.device = device
+
+    @staticmethod
+    def simple_dqn(model, **kwargs):
+        return QLearningPreprocessor(model=model, target_model=None, use_double_dqn=False, **kwargs)
+
+    @staticmethod
+    def target_dqn(model, target_model, **kwards):
+        return QLearningPreprocessor(model, target_model, use_double_dqn=False, **kwards)
+
+    @staticmethod
+    def double_dqn(model, target_model, **kwargs):
+        return QLearningPreprocessor(model, target_model, use_double_dqn=True, **kwargs)
+
+    def _calc_Q(self, states_first, states_last):
+        """
+        Calculates apropriate q values for first and last states. Way of calculate depends on our settings.
+        :param states_first: numpy array of first states
+        :param states_last: numpy array of last states
+        :return: tuple of numpy arrays of q values
+        """
+        # here we need both first and last values calculated using our main model, so we
+        # combine both states into one batch for efficiency and separate results later
+        if self.target_model is None or self.use_double_dqn:
+            states_t = torch.tensor(np.concatenate((states_first, states_last), axis=0)).to(self.device)
+            res_both = self.model(states_t).data.cpu().numpy()
+            return res_both[:len(states_first)], res_both[len(states_first):]
+
+        # in this case we have target_model set and use_double_dqn==False
+        # so, we should calculate first_q and last_q using different models
+        states_first_v = torch.tensor(states_first).to(self.device)
+        states_last_v = torch.tensor(states_last).to(self.device)
+        q_first = self.model(states_first_v).data
+        q_last = self.target_model(states_last_v).data
+        return q_first.cpu().numpy(), q_last.cpu().numpy()
+
+    def _calc_target_rewards(self, states_last, q_last):
+        """
+        Calculate rewards from final states according to variants from our construction:
+        1. simple DQN: max(Q(states, model))
+        2. target DQN: max(Q(states, target_model))
+        3. double DQN: Q(states, target_model)[argmax(Q(states, model)]
+        :param states_last: numpy array of last states from the games
+        :param q_last: numpy array of last q values
+        :return: vector of target rewards
+        """
+        # in this case we handle both simple DQN and target DQN
+        if self.target_model is None or not self.use_double_dqn:
+            return q_last.max(axis=1)
+
+        # here we have target_model set and use_double_dqn==True
+        actions = q_last.argmax(axis=1)
+        # calculate Q values using target net
+        states_last_v = torch.tensor(states_last).to(self.device)
+        q_last_target = self.target_model(states_last_v).data.cpu().numpy()
+        return q_last_target[range(q_last_target.shape[0]), actions]
+
+    def preprocess(self, batch):
+        """
+        Calculates data for Q learning from batch of observations
+        :param batch: list of lists of Experience objects
+        :return: tuple of numpy arrays:
+            1. states -- observations
+            2. target Q-values
+            3. vector of td errors for every batch entry
+        """
+        # first and last states for every entry
+        state_0 = np.array([exp[0].state for exp in batch], dtype=np.float32)
+        state_L = np.array([exp[-1].state for exp in batch], dtype=np.float32)
+
+        q0, qL = self._calc_Q(state_0, state_L)
+        rewards = self._calc_target_rewards(state_L, qL)
+
+        td = np.zeros(shape=(len(batch),))
+
+        for idx, (total_reward, exps) in enumerate(zip(rewards, batch)):
+            # game is done, no final reward
+            if exps[-1].done:
+                total_reward = 0.0
+            for exp in reversed(exps[:-1]):
+                total_reward *= self.gamma
+                total_reward += exp.reward
+            # update total reward and calculate td error
+            act = exps[0].action
+            td[idx] = q0[idx][act] - total_reward
+            q0[idx][act] = total_reward
+
+        return state_0, q0, td

simonbb-0.0.4/src/simonbb/cis433/ignite.py

@@ -0,0 +1,137 @@
+import time
+from typing import Optional
+from ignite.engine import Engine, State
+from ignite.engine import Events, EventEnum
+from ignite.handlers.timing import Timer
+
+from .experience import ExperienceSource
+
+class EpisodeEvents(EventEnum):
+    EPISODE_COMPLETED = "episode_completed"
+    BOUND_REWARD_REACHED = "bound_reward_reached"
+    BEST_REWARD_REACHED = "best_reward_reached"
+
+
+class EndOfEpisodeHandler:
+    def __init__(self, exp_source: ExperienceSource, alpha: float = 0.98,
+                 bound_avg_reward: Optional[float] = None,
+                 subsample_end_of_episode: Optional[int] = None):
+        """
+        Construct end-of-episode event handler
+        :param exp_source: experience source to use
+        :param alpha: smoothing alpha param
+        :param bound_avg_reward: optional boundary for average reward
+        :param subsample_end_of_episode: if given, end of episode event will be subsampled by this amount
+        """
+        self._exp_source = exp_source
+        self._alpha = alpha
+        self._bound_avg_reward = bound_avg_reward
+        self._best_avg_reward = None
+        self._subsample_end_of_episode = subsample_end_of_episode
+
+    def attach(self, engine: Engine):
+        engine.add_event_handler(Events.ITERATION_COMPLETED, self)
+        engine.register_events(*EpisodeEvents)
+        State.event_to_attr[EpisodeEvents.EPISODE_COMPLETED] = "episode"
+        State.event_to_attr[EpisodeEvents.BOUND_REWARD_REACHED] = "episode"
+        State.event_to_attr[EpisodeEvents.BEST_REWARD_REACHED] = "episode"
+
+    def __call__(self, engine: Engine):
+        for reward, steps in self._exp_source.pop_rewards_steps():
+            engine.state.episode = getattr(engine.state, "episode", 0) + 1
+            engine.state.episode_reward = reward
+            engine.state.episode_steps = steps
+            engine.state.metrics['reward'] = reward
+            engine.state.metrics['steps'] = steps
+            self._update_smoothed_metrics(engine, reward, steps)
+            if self._subsample_end_of_episode is None or engine.state.episode % self._subsample_end_of_episode == 0:
+                engine.fire_event(EpisodeEvents.EPISODE_COMPLETED)
+            if self._bound_avg_reward is not None and engine.state.metrics['avg_reward'] >= self._bound_avg_reward:
+                engine.fire_event(EpisodeEvents.BOUND_REWARD_REACHED)
+            if self._best_avg_reward is None:
+                self._best_avg_reward = engine.state.metrics['avg_reward']
+            elif self._best_avg_reward < engine.state.metrics['avg_reward']:
+                engine.fire_event(EpisodeEvents.BEST_REWARD_REACHED)
+                self._best_avg_reward = engine.state.metrics['avg_reward']
+
+    def _update_smoothed_metrics(self, engine: Engine, reward: float, steps: int):
+        for attr_name, val in zip(('avg_reward', 'avg_steps'), (reward, steps)):
+            if attr_name not in engine.state.metrics:
+                engine.state.metrics[attr_name] = val
+            else:
+                engine.state.metrics[attr_name] *= self._alpha
+                engine.state.metrics[attr_name] += (1-self._alpha) * val
+
+
+class EpisodeFPSHandler:
+    FPS_METRIC = 'fps'
+    AVG_FPS_METRIC = 'avg_fps'
+    TIME_PASSED_METRIC = 'time_passed'
+
+    def __init__(self, fps_mul: float = 1.0, fps_smooth_alpha: float = 0.98):
+        self._timer = Timer(average=True)
+        self._fps_mul = fps_mul
+        self._started_ts = time.time()
+        self._fps_smooth_alpha = fps_smooth_alpha
+
+    def attach(self, engine: Engine, manual_step: bool = False):
+        self._timer.attach(engine, step=None if manual_step else Events.ITERATION_COMPLETED)
+        engine.add_event_handler(EpisodeEvents.EPISODE_COMPLETED, self)
+        engine.state.metrics[self.AVG_FPS_METRIC] = 0
+
+    def step(self):
+        """
+        If manual_step=True on attach(), this method should be used every time we've communicated with environment
+        to get proper FPS
+        :return:
+        """
+        self._timer.step()
+
+    def __call__(self, engine: Engine):
+        t_val = self._timer.value()
+        if engine.state.iteration > 1:
+            fps = self._fps_mul / t_val
+            avg_fps = engine.state.metrics.get(self.AVG_FPS_METRIC)
+            if avg_fps is None or avg_fps <= 0:
+                avg_fps = fps
+            else:
+                avg_fps *= self._fps_smooth_alpha
+                avg_fps += (1-self._fps_smooth_alpha) * fps
+            engine.state.metrics[self.AVG_FPS_METRIC] = avg_fps
+            engine.state.metrics[self.FPS_METRIC] = fps
+        engine.state.metrics[self.TIME_PASSED_METRIC] = time.time() - self._started_ts
+        self._timer.reset()
+
+
+class PeriodEvents(EventEnum):
+    ITERS_10_COMPLETED = "iterations_10_completed"
+    ITERS_100_COMPLETED = "iterations_100_completed"
+    ITERS_1000_COMPLETED = "iterations_1000_completed"
+    ITERS_10000_COMPLETED = "iterations_10000_completed"
+    ITERS_100000_COMPLETED = "iterations_100000_completed"
+
+
+class PeriodicEvents:
+    """
+    The same as CustomPeriodicEvent from ignite.contrib, but use true amount of iterations,
+    which is good for TensorBoard
+    """
+
+    INTERVAL_TO_EVENT = {
+        10: PeriodEvents.ITERS_10_COMPLETED,
+        100: PeriodEvents.ITERS_100_COMPLETED,
+        1000: PeriodEvents.ITERS_1000_COMPLETED,
+        10000: PeriodEvents.ITERS_10000_COMPLETED,
+        100000: PeriodEvents.ITERS_100000_COMPLETED,
+    }
+
+    def attach(self, engine: Engine):
+        engine.add_event_handler(Events.ITERATION_COMPLETED, self)
+        engine.register_events(*PeriodEvents)
+        for e in PeriodEvents:
+            State.event_to_attr[e] = "iteration"
+
+    def __call__(self, engine: Engine):
+        for period, event in self.INTERVAL_TO_EVENT.items():
+            if engine.state.iteration % period == 0:
+                engine.fire_event(event)