import copy
import logging
import multiprocessing
import time
from collections import defaultdict
import math
import scipy.stats
import gymnasium as gym
import numpy as np
from verigym.environments.interval_explicitenv import IntervalAbstractedExplicitEnv
from ..environments.reward_func import RewardFunction, IntervalRewardFunction
from ..environments.transition_func import TransitionFunction, IntervalTransitionFunction
from ..environments.state_distributions import StateDistribution
from ..environments.explicitenv import AbstractedExplicitEnv
from ..environments.verigymenv import VeriGymEnv
from ..environments.labeling import AbstractStateLabeler
from ..policy.policy import PolicyClass
from .abstractionmapper import AbstractionMapper, validate_for_abstraction
logger = logging.getLogger(__name__)
[docs]
def get_interval_transition_reward(
T_counts, R_cumsum, P_tot_counts, n_states, n_actions, confidence=0.9, iid=False
):
"""
Construct (confidence) intervals from the observed data.
If iid, computes Clopper-Pearson Binomial intervals.
Else, computes Azuma-Hoeffding Martingale intervals.
Expects the raw counts (see `learn_abstraction()`), not the normalized probabilities.
Intervals are only constructed for observed successors; unobserved successors keep the interval (0.0, 0.0),
i.e., the same support as the point estimates of `normalize_aggregated_counts()`.
Parameters
----------
T_counts : dict
Mapping s -> a -> s' -> number of observed transitions (s, a, s')
R_cumsum : dict
Mapping s -> a -> cumulative observed reward (the mean is R_cumsum[s][a] / P_tot_counts[(s, a)])
P_tot_counts : dict
Mapping of (s,a) to total counts (e.g., the sum of successor counts in T_counts for each s,a).
State-action pairs with a total count of 0 are skipped.
n_states : int
``|S|``
n_actions : int
``|A|``
confidence : float, optional
The statistical (high) confidence of the iMDP construction, by default 0.9
iid : boolean, optional
Whether the data is independently and identically distributed (i.i.d.), by default False
"""
# TODO: implement reward estimation from IID data.
delta = 1 - confidence # Confidence over total model
# M = len(P_tot_counts) * n_states # TODO: Use actual visited counts for the confidence guarantee instead of whole state space?
M = n_states**2 * n_actions
alpha = delta / M # confidence for each transition
interval_T = make_interval_transition_dict()
interval_R = make_reward_dict()
for (s, a), n in P_tot_counts.items():
if n == 0: # unvisited, skipped just like in `normalize_aggregated_counts()`
continue
for ss, k in T_counts[s][a].items():
if not isinstance(k, (int, np.integer)):
raise TypeError(
f"Expected raw (integer) counts, but T_counts[{s}][{a}][{ss}] = {k!r}. "
"Intervals must be computed from counts, not from normalized probabilities."
)
if iid: # Use Clopper-Pearson Binomial intervals
lb = 0.0 if k == 0 else scipy.stats.beta.ppf(alpha / 2, k, n - k + 1)
ub = 1.0 if k == n else scipy.stats.beta.ppf(1 - alpha / 2, k + 1, n - k)
else: # Use Azuma-Hoeffding Martingale intervals
eps = math.sqrt(math.log(2 / alpha) / (2 * n))
lb = 0.0 if k == 0 else k / n - eps
ub = 1.0 if k == n else k / n + eps
lb = max(0.0, lb)
ub = min(1.0, ub)
interval_T[s][a][ss] = (lb, ub)
mean_reward = R_cumsum[s][a] / n
interval_R[s][a] = (mean_reward, mean_reward) # FIXME? Use expected / MLE for now
interval_T_function = IntervalTransitionFunction(n_states, n_actions, interval_T)
interval_R_function = IntervalRewardFunction(n_states, n_actions, interval_R)
return interval_T_function, interval_R_function
[docs]
def create_abstraction(
original_env: VeriGymEnv,
abstraction_mapper: AbstractionMapper,
exploration_policy: PolicyClass,
num_steps: int,
n_iterations: int = 1,
multithreading: bool = True,
verbose: bool = False,
return_info = False,
intervals: bool = False,
assume_iid: bool = False,
dataset: list[list[tuple]] | None = None,
seed: int = None
) -> AbstractedExplicitEnv:
"""
Creates an abstraction from a VeriGymEnv by discretizing the state and
action spaces. Returns an `AbstractedExplicitEnv`.
Parameters
----------
original_env : VeriGymEnv
The environment / model to be abstracted.
abstraction_mapper: AbstractionMapper
Holds the mapping between the original space (of `original_env`) and the space which we are abstracting into.
exploration_policy : PolicyClass
The policy of interacting with the `original_env` (e.g. a random policy).
num_steps : int
Number of steps to simulate within the environment per iteration (also, see `n_iterations` and `dataset`).
Set to 0 to only use a given `dataset` without any simulation.
n_iterations: int
Number of (interleaving) iterations. For each iteration the `exploration_policy.update_for_abstraction_refinement(...)`
will be called, allowing the policy to update based on the gathered interactions (e.g. for state-based exploration
policies). Note, that the total number of simulated steps equals `n_iterations * num_steps`, or
`(n_iterations - 1) * num_steps` if a `dataset` is given. For policies that are not interleaving, set `n_iterations`
to 1 (default).
multithreading: bool, optional
Whether to multithread or use single thread.
verbose : bool, optional
Whether to be verbose, by default False.
intervals: bool, optional
Whether to learn transition functions with (confidence) intervals or point-based estimates.
Default is False, i.e., point-based estimate learning.
assume_iid: bool, optional.
Whether to assume that the data distribution is iid or not. This is only relevant for interval learning.
Defaults to False.
dataset: list[list[tuple]] | None, optional
Pre-collected trajectories (as returned by `VeriGymEnv.simulate()`) in the original space. If given,
the dataset is used in the first of `n_iterations` and only the following iterations collect the `num_steps` samples.
To only use the dataset without any simulation, set `n_iterations == 1` and `num_steps == 0`. Defaults to None.
seed: int | None, optional
Seeds the simulation of `original_env` (a different seed per iteration, derived from `seed`) and the
`exploration_policy` (with `exploration_policy.seed(seed)`, which overrides the policy's own seed).
If None (default), both continue with their current random generators.
Returns
-------
AbstractedExplicitEnv
The abstracted model.
Notes
-----
- If intervals == True, the return type is `IntervalAbstractedExplicitEnv`, which is a sub-type of `AbstractedExplicitEnv`.
It only stores the intervals, and its point estimates (e.g., to sample transitions) are derived from them,
see `IntervalTransitionFunction.point_estimate`.
- The returned model has an extra terminal sink state with index `abstraction_mapper.abstract_n_states` (the last
enumerated state). Hence it has `abstract_n_states + 1` states and its `observation_space` is
`Discrete(abstract_n_states + 1)`.
- Every step on which `original_env` terminated (not truncated) transitions into the terminal state. The abstraction mapper never maps an observation to it (meaning no original state maps to the terminal state in the abstract state space).
- Every action of the terminal state loops back to itself with reward 0 (see `add_terminal_sink_state()`). Without them, the
terminal state would have no actions, i.e., it would be a deadlock.
- In interval models, these self-loops are given the exact interval (1.0, 1.0).
"""
assert isinstance(original_env, gym.Env), (
f"original_env is type {type(original_env)} and does not inherit from gym.Env"
)
assert isinstance(n_iterations, int) and n_iterations > 0, f"n_iterations must be an int larger than 0, {n_iterations = }"
validate_for_abstraction(abstraction_mapper, multithreading=multithreading)
assert dataset is None or n_iterations > 1 or num_steps == 0, (
f"A given dataset together with n_iterations==1 but num_steps>0, is not compatible. The given num_steps {num_steps = } would be ignored. Set num_steps=0 to only use the dataset, or n_iterations > 1 to simulate on top of it."
)
if intervals:
if assume_iid:
if (
abstraction_mapper.original_n_states == math.inf
or abstraction_mapper.original_n_actions == math.inf
):
print(
f"WARNING: Can only learn intervals from IID data on discrete environments, but received: n_states={abstraction_mapper.original_n_states} and n_actions={abstraction_mapper.original_n_actions}."
)
else:
if (
abstraction_mapper._state_abstraction_map.is_identity_map
and abstraction_mapper._action_abstraction_map.is_identity_map
):
print(
"WARNING: Data said to be non-IID when constructing invervals but creating abstraction through identity mapping."
)
# Get discrete states and actions
n_states = abstraction_mapper.abstract_n_states
n_actions = abstraction_mapper.abstract_n_actions
assert (n_states is not None) and (n_actions is not None), (
f"Neither should be none {(n_states, n_actions) = }"
)
# cast from np.int64, as the reward function requires python ints as state indices
n_states, n_actions = int(n_states), int(n_actions)
# add terminal state: Terminated steps transition into an extra absorbing state behind the states of the mapper
terminal_state = n_states
n_states += 1
# Initialize relevant objects for learning the abstraction
T_counts, R_cumsum, P_tot_counts, state_distr_counts = _create_count_databases()
iteration_dataset = []
# Initialize info dict
info = {
"time_simulate":0.0,
"time_policy_update":0.0,
"time_model_construction":0.0
}
# Loop through iterations. If interleaving abstraction is not required, n_iterations will be just 1.
if seed is not None:
seeds_list = np.random.SeedSequence(seed).generate_state(n_iterations).tolist() # use different seeds at each iteration
exploration_policy.seed(seed)
else:
seeds_list = [None] * n_iterations # continue with the current random state of the env
for i in range(n_iterations):
t0 = time.time()
exploration_policy = exploration_policy.update_for_abstraction_refinement(
iteration_dataset, T_counts, P_tot_counts, R_cumsum, state_distr_counts
)
assert isinstance(exploration_policy, PolicyClass)
t1 = time.time()
t_policy = t1 - t0
print(f"Policy update time: {t_policy:.4f}s") if verbose else {}
info["time_policy_update"] += t_policy
if dataset is not None and i == 0:
# if we already have a starting dataset, in the very first iteration we use it
iteration_dataset = dataset
else:
# generate dataset via simulation
iteration_dataset = original_env.simulate( # TODO get rid of this simulate call, is it requires original_env to be of type VeriGymEnv. This should also work for gym.Env
policy=exploration_policy, n_steps=num_steps, verbose=verbose, seed=seeds_list[i]
)
t2 = time.time()
t_simulate = t2 - t1
print(f"Simulation time: {t_simulate:.4f}s") if verbose else {}
info["time_simulate"] += t_simulate
# approximate the transition function from new dataset
# note, we are only getting the counts for state/action/nex_state/reward pairs here
new_T_counts, new_R_cumsum, new_P_tot_counts, new_state_distr_counts = (
learn_abstraction(
iteration_dataset,
n_states,
n_actions,
abstraction_mapper=abstraction_mapper,
multithreading=multithreading,
terminal_state=terminal_state,
verbose=verbose
)
)
# --- START Aggregate across iterations
for s, count in new_state_distr_counts.items():
state_distr_counts[s] += count
for (s, a), tot_count in new_P_tot_counts.items():
P_tot_counts[(s, a)] += tot_count
R_cumsum[s][a] += new_R_cumsum[s][a]
for next_state, count in new_T_counts[s][a].items():
T_counts[s][a][next_state] += count
# --- END Aggregate
t3 = time.time()
t_construct = t3 - t2
print(f"Learning Abstraction: {t_construct:.4f}s") if verbose else {}
info["time_model_construction"] += t_construct
t0 = time.time()
# add terminal/sink state
add_terminal_sink_state(T_counts, R_cumsum, P_tot_counts, terminal_state, n_actions)
# Obtain valid distributions/values by aggregating the variables storing the counts (normalizing via P_tot_counts)
T, R, S_init = normalize_aggregated_counts(
T_counts, R_cumsum, P_tot_counts, state_distr_counts, n_states, n_actions
)
info["time_model_construction"] += time.time() - t0
shared_args = {
"nr_states" : n_states,
"nr_actions" : n_actions,
"nr_rewards" : 1,
"initial_state_distr" : S_init,
"abstraction_map" : abstraction_mapper,
"original_env" : original_env,
"render_mode" : None,
"dummy_terminal_state" : terminal_state
}
if intervals:
interval_T, interval_R = get_interval_transition_reward(
T_counts, R_cumsum, P_tot_counts, n_states, n_actions, iid=assume_iid
)
# the self-loops of the terminal sink state are known, not estimated
for a in range(n_actions):
interval_T.T_dict[terminal_state][a][terminal_state] = (1.0, 1.0)
abstracted_env = IntervalAbstractedExplicitEnv(
**shared_args,
interval_transition_function=interval_T,
interval_reward_function=interval_R,
)
else:
# Construct the abstracted ExplicitEnv
abstracted_env = AbstractedExplicitEnv(
**shared_args,
transition_function=T,
reward_function=R,
)
abstracted_env.state_labeler = AbstractStateLabeler(
original_labeler=original_env.state_labeler,
abstraction_mapper=abstraction_mapper,
)
if return_info:
return abstracted_env, info
return abstracted_env
def _create_count_databases() -> tuple[dict, dict, dict, dict]:
"""
Creates all required objects for the abstraction learning.
The returned objects `T_counts`, `R_cumsum` and `state_distr_counts` have the same datatype
as their usual counterparts, only that they are intended to hold the absolute number of samples
(or, for `R_cumsum`, the summed up reward of all dataset samples corresponding to a state-action pair) and are not normalized to probability distributions /
mean rewards. For obtaining probablities one has to normalize using the total counts stored in `P_tot_counts`.
This `dict` is populated for convenience as its usage reduces the amount of times the counts
would have to be computed from `T_counts`.
Parameters
----------
n_states : int
The number of states.
Returns
-------
tuple[dict, dict, dict, dict]
T_counts, R_cumsum, P_tot_counts, state_distr_counts.
Example
-------
::
import math
from verigym.abstraction.learn_abstraction import _create_count_databases
bin_edges_states = [[2, 3], [0.5, 1.0, 1.5]]
bin_edges_actions = [[-0.5, 0.0, 0.5]]
# number of states
n_states = math.prod([len(dimension) - 1 for dimension in bin_edges_states])
# number of counts (occurences) for each state-action-next_state pair
(
T_counts,
R_cumsum,
P_tot_counts,
state_distr_counts,
) = _create_count_databases(n_states)
"""
T_counts = make_transition_dict()
# cumulative reward for all individual state-action pairs (not cumulative over time)
R_cumsum = make_reward_cumsum_dict()
# dict with occurences for each state-action pair
P_tot_counts = make_int_dict()
# dict with occurences of each state as initial state
state_distr_counts = make_int_dict()
return (
T_counts,
R_cumsum,
P_tot_counts,
state_distr_counts,
)
# Define named functions for defaultdict factories (they cannot be lambda functions due to multiprocess pickling issues)
def make_int_dict(): # pragma: no cover
return defaultdict(int) # calling int() without argument returns 0
def make_list_dict(): # pragma: no cover
return defaultdict(list)
def make_middle_dict(): # pragma: no cover
return defaultdict(make_int_dict)
def make_transition_dict():
"""Dict that can be used for transition function"""
return defaultdict(make_middle_dict)
def make_zero_interval(): # pragma: no cover
return (0.0, 0.0)
def make_interval_dict(): # pragma: no cover
return defaultdict(make_zero_interval)
def make_interval_middle_dict(): # pragma: no cover
return defaultdict(make_interval_dict)
def make_interval_transition_dict():
"""Dict that can be used for interval transition function"""
return defaultdict(make_interval_middle_dict)
def make_reward_dict():
"""Dict that can be used for reward function"""
return defaultdict(make_list_dict)
def make_float_dict(): # pragma: no cover
return defaultdict(float) # calling float() without argument returns 0.0
def make_reward_cumsum_dict():
"""Dict that can be used for accumulating rewards"""
return defaultdict(make_float_dict)
def add_terminal_sink_state(
T_counts: dict, R_cumsum: dict, P_tot_counts: dict, terminal_state_enum: int, n_actions: int
) -> None:
"""
Adds outgoing transitions of the terminal sink state to the count/reward databases (`T_counts`, `R_cumsum`, `P_tot_counts`), in place.
For every action, the terminal state transitions to itself with reward 0, so once it is entered, it is never
left. Call this once on the aggregated counts, before normalizing them (see `normalize_aggregated_counts()`).
Parameters
----------
T_counts : dict
Mapping s -> a -> s' -> number of observed transitions (s, a, s').
R_cumsum : dict
Mapping s -> a -> sum of observed rewards.
P_tot_counts : dict
Mapping of (s, a) to total counts.
terminal_state_enum : int
Enumerated index of the terminal sink state.
n_actions : int
Number of actions.
"""
for a in range(n_actions):
T_counts[terminal_state_enum][a][terminal_state_enum] += 1
P_tot_counts[(terminal_state_enum, a)] += 1
R_cumsum[terminal_state_enum][a] = 0.0
[docs]
def collect_data_from_trajectories(
trajectories: list[list[tuple[int, int, float, int, bool]]],
n_states: int,
mapper: AbstractionMapper | None = None,
terminal_state: int | None = None,
) -> tuple[dict, dict, dict, dict]:
"""
Taking a dataset of `trajectories` populates dicts counting the total
occurences in the `trajectories` and add them to the objects that can then
be used for computing the transition and reward function as well as the
initial state distribution.
Parameters
----------
trajectories : list[list[tuple[int, int, float, int, bool]]]
Dataset of trajectories (state, action, reward, next_state, terminated), as returned by
`VeriGymEnv.simulate()`. Steps without the `terminated` flag, i.e. (state, action, reward, next_state),
are treated as not terminated.
n_states : int
Number of states of the corresponding state space (including `terminal_state`, if given).
mapper : AbstractionMapper | None
Maps from the state and action spaces of an original environment to an abstract environment.
If None, an identity map (no mapping) will be perormed. Defaults to `None`.
terminal_state : int | None
State that terminated steps transition into, instead of the (mapped) `next_state`.
If None, terminated steps are counted like any other step. Defaults to `None`.
Returns
-------
tuple[dict, dict, dict, dict]
T_counts, R_cumsum, P_tot_counts, state_distr_counts. The count
databases created by `_create_count_databases`, populated from `trajectories`.
"""
# If no mapper is passed, we keep states and actions as they are
if mapper is None:
def original_to_abstract_state(x):
return x
def original_to_abstract_action(x):
return x
else:
# using `*_enum` functions.
original_to_abstract_state = mapper._state_abstraction_map.original_to_enum
original_to_abstract_action = mapper._action_abstraction_map.original_to_enum
assert (original_to_abstract_state is not None) and (
original_to_abstract_action is not None
), (
f"One of the abstraction maps is None: {original_to_abstract_state = }, {original_to_abstract_action = }"
)
# Initialize local storage for this thread
(
T_counts,
R_cumsum,
P_tot_counts,
state_distr_counts,
) = _create_count_databases()
for trajectory in trajectories:
for i, (s, a, r, s_next, *terminated) in enumerate(trajectory):
if isinstance(r, np.ndarray):
r = r.item()
# Go through the mapper wrappers (not `.forward_map` directly) so
# that size-1 ndarray outputs are normalized to hashable scalars,
# which is required for use as dict keys / array indices below.
s = original_to_abstract_state(s)
a = original_to_abstract_action(a)
if terminal_state is not None and terminated and terminated[0]:
s_next = terminal_state
else:
s_next = original_to_abstract_state(s_next)
if i == 0:
state_distr_counts[s] += 1
T_counts[s][a][s_next] += 1
P_tot_counts[(s, a)] += 1
R_cumsum[s][a] += float(r)
return T_counts, R_cumsum, P_tot_counts, state_distr_counts
def learn_abstraction_multithreaded(
dataset: list[list[tuple[int, int, float, int, bool]]],
n_states: int,
n_actions: int,
abstraction_mapper: AbstractionMapper,
terminal_state: int | None = None,
verbose = True
):
num_threads = max(min(4, multiprocessing.cpu_count() - 1), 1)
chunk_size = len(dataset) // num_threads
# One worker (few cores or a tiny dataset) would only add process overhead
if num_threads == 1 or chunk_size == 0:
return learn_abstraction_single_threaded(
dataset, n_states, n_actions, abstraction_mapper, terminal_state
)
(
T_dict,
R_dict,
P_tot,
state_distr,
) = _create_count_databases()
chunks = [
(
dataset[i * chunk_size : i * chunk_size + chunk_size],
n_states,
copy.deepcopy(abstraction_mapper),
terminal_state,
)
for i in range(num_threads - 1)
]
chunks.append(
(
dataset[(num_threads - 1) * chunk_size :],
n_states,
copy.deepcopy(abstraction_mapper),
terminal_state,
)
)
lens = [len(chunk[0]) for chunk in chunks]
assert sum(lens) == len(dataset), f"{sum(lens)=} and {len(dataset)=}"
tik = time.time()
# forkserver: forking a process with running threads can deadlock the workers.
# close() + join() lets the workers exit on their own; otherwise `__exit__` terminates
# them with SIGTERM, which can hang them
with multiprocessing.get_context("forkserver").Pool(num_threads) as executor:
results = executor.starmap(collect_data_from_trajectories, chunks)
executor.close()
executor.join()
tok = time.time()
print("processing in ", tok - tik) if verbose else {}
print("aggregating..") if verbose else {}
for _T_results, _R_results, P_tot_results, state_distr_results in results:
for s, count in state_distr_results.items():
state_distr[s] += count
for (s, a), tot_count in P_tot_results.items():
P_tot[(s, a)] += tot_count
# results are unpacked once rather than once per state-action pair.
all_T_results = [
T_results
for T_results, _R_results, _P_tot_results, _state_distr_results in results
]
for (s, a), tot_count in P_tot.items():
if tot_count == 0:
continue
for T_results in all_T_results:
if s in T_results and a in T_results[s]:
for s_next, count in T_results[s][a].items():
T_dict[s][a][s_next] += count
for _T_results, R_results, _P_tot_results, _state_distr_results in results:
for s in R_results:
for a in R_results[s]:
R_dict[s][a] += R_results[s][a]
print("aggregating in", time.time() - tok) if verbose else {}
return T_dict, R_dict, P_tot, state_distr
[docs]
def learn_abstraction(
dataset: list[list[tuple[int, int, float, int, bool]]],
n_states: int,
n_actions: int,
abstraction_mapper: AbstractionMapper = None,
multithreading: bool = True,
terminal_state: int | None = None,
verbose=True
) -> tuple[dict, dict, dict, dict]:
"""
Abstraction learning for a given dataset. Single- or multithreaded.
Computes the total counts (!) for transition and reward function and initial state distribution.
Do not forget to normalize (see `normalize_aggregated_counts()`) for obtaining probability distributions.
Parameters
----------
dataset : list[list[tuple[int, int, float, int, bool]]]
The dataset of which we are learning the abstraction, with steps (state, action, reward, next_state, terminated).
n_states : int
Number of states in the state space. (corresponds to the abstract state space if `abstraction_mapper` is not `None`)
Includes `terminal_state`, if given.
n_actions : int
Number of actions in the action space. (corresponds to the abstract action space if `abstraction_mapper` is not `None`)
abstraction_mapper : AbstractionMapper, optional
Mapping from an original space (samples in dataset) to the abstract space. If no mapping is required, set to `None`, by default None
multithreading : bool, optional
Flag for using single- or multithreading, by default True
terminal_state : int | None, optional
State (number) that terminated steps transition to (see `collect_data_from_trajectories()`).
The terminal's self-loops are not added here, see `add_terminal_sink_state()`. By default None
Returns
-------
tuple[dict, dict, dict, dict]
T_counts, R_cumsum, P_tot_counts, state_distr_counts
"""
print(f"Trajectories in dataset: {len(dataset)}") if verbose else {}
if multithreading:
return learn_abstraction_multithreaded(
dataset, n_states, n_actions, abstraction_mapper, terminal_state, verbose = verbose
)
else:
return learn_abstraction_single_threaded(
dataset, n_states, n_actions, abstraction_mapper, terminal_state
)
[docs]
def normalize_aggregated_counts(
T_dict, R_dict, P_tot, state_distr, n_states, n_actions
):
"""
Normalizes the aggregated counts (see `learn_abstraction()`) into point estimates.
The inputs are not modified in place.
Parameters
----------
T_dict : dict
Mapping s -> a -> s' -> number of observed transitions (s, a, s').
R_dict : dict
Mapping s -> a -> cumulative reward for each state-action pair (not over time).
P_tot : dict
Mapping of (s, a) to total counts. State-action pairs with a total count of 0 are skipped.
state_distr : dict
Mapping s -> number of occurences of s as initial state.
n_states : int
Number of states.
n_actions : int
Number of actions.
Returns
-------
tuple[TransitionFunction, RewardFunction, StateDistribution]
T, R, S_init
"""
tot_init_count = sum(state_distr.values())
S_init = StateDistribution(
n_states, {s: count / tot_init_count for s, count in state_distr.items() if count > 0}
)
T_normalized = make_transition_dict()
R_normalized = make_reward_dict()
for (s, a), tot_count in P_tot.items():
if tot_count == 0:
continue
for s_next, count in T_dict[s][a].items():
T_normalized[s][a][s_next] = count / tot_count
sum_tot = sum(T_normalized[s][a].values())
assert round(sum_tot, 1) in {0, 1}, (
f"Counts for {s, a} sum to {sum_tot} != {0, 1}!"
)
R_normalized[s][a] = R_dict[s][a] / tot_count
return (
TransitionFunction(n_states, n_actions, T_normalized),
RewardFunction(n_states, n_actions, R_normalized),
S_init,
)
def learn_abstraction_single_threaded(
dataset: list[list[tuple[int, int, float, int, bool]]],
n_states: int,
n_actions: int,
abstraction_mapper: AbstractionMapper,
terminal_state: int | None = None,
) -> tuple[dict, dict, dict, dict]:
T_counts, R_cumsum, P_tot_counts, state_distr_counts = (
collect_data_from_trajectories(dataset, n_states, abstraction_mapper, terminal_state)
)
return T_counts, R_cumsum, P_tot_counts, state_distr_counts