def __init__(self, nb_states, obs_dim, act_dim,
nb_lags, prior, degree=1, likelihood=None):
assert nb_lags > 0
self.nb_states = nb_states
self.obs_dim = obs_dim
self.act_dim = act_dim
self.nb_lags = nb_lags
self.degree = degree
self.feat_dim = int(sc.special.comb(self.degree + self.obs_dim, self.degree)) - 1
self.basis = PolynomialFeatures(self.degree, include_bias=False)
self.input_dim = self.feat_dim + 1
self.output_dim = self.act_dim
self.prior = prior
self.posterior = copy.deepcopy(prior)
# Linear-Gaussian likelihood
if likelihood is not None:
self.likelihood = likelihood
else:
As, lmbdas = self.prior.rvs()
self.likelihood = StackedLinearGaussiansWithPrecision(size=self.nb_states,
column_dim=self.input_dim,
row_dim=self.output_dim,
As=As, lmbdas=lmbdas, affine=True)
def __init__(self, nb_states, ltn_dim, act_dim,
nb_lags, prior, likelihood=None):
super(BayesianAutoRegressiveGaussianLatent, self).__init__(nb_states, ltn_dim, act_dim,
nb_lags, prior, likelihood)
# Linear-Gaussian likelihood
if likelihood is not None:
self.likelihood = likelihood
else:
As, lmbdas = self.prior.rvs()
self.likelihood = StackedLinearGaussiansWithPrecision(size=self.nb_states,
column_dim=self.input_dim,
row_dim=self.output_dim,
As=As, lmbdas=lmbdas, affine=True)
class BayesianInitGaussianControl:
def __init__(self, nb_states, obs_dim, act_dim,
nb_lags, prior, degree=1, likelihood=None):
assert nb_lags > 0
self.nb_states = nb_states
self.obs_dim = obs_dim
self.act_dim = act_dim
self.nb_lags = nb_lags
self.degree = degree
self.feat_dim = int(sc.special.comb(self.degree + self.obs_dim, self.degree)) - 1
self.basis = PolynomialFeatures(self.degree, include_bias=False)
self.input_dim = self.feat_dim + 1
self.output_dim = self.act_dim
self.prior = prior
self.posterior = copy.deepcopy(prior)
# Linear-Gaussian likelihood
if likelihood is not None:
self.likelihood = likelihood
else:
As, lmbdas = self.prior.rvs()
self.likelihood = StackedLinearGaussiansWithPrecision(size=self.nb_states,
column_dim=self.input_dim,
row_dim=self.output_dim,
As=As, lmbdas=lmbdas, affine=True)
@property
def params(self):
return self.likelihood.params
@params.setter
def params(self, values):
self.likelihood.params = values
def permute(self, perm):
self.likelihood.As = self.likelihood.As[perm]
self.likelihood.lmbdas = self.likelihood.lmbdas[perm]
def initialize(self, x, u, **kwargs):
kmeans = kwargs.get('kmeans', False)
x0, u0 = [], []
for _x, _u in zip(x, u):
x0.append(_x[:self.nb_lags])
u0.append(_u[:self.nb_lags])
f0 = list(map(self.featurize, x0))
t = list(map(len, f0))
if kmeans:
from sklearn.cluster import KMeans
km = KMeans(self.nb_states)
km.fit(np.vstack(f0))
z0 = np.split(km.labels_, np.cumsum(t)[:-1])
else:
z0 = list(map(partial(npr.choice, self.nb_states), t))
z0 = list(map(partial(one_hot, self.nb_states), z0))
stats = self.likelihood.weighted_statistics(f0, u0, z0)
self.posterior.nat_param = self.prior.nat_param + stats
self.likelihood.params = self.posterior.rvs()
def featurize(self, x):
feat = self.basis.fit_transform(np.atleast_2d(x))
return np.squeeze(feat) if x.ndim == 1\
else np.reshape(feat, (x.shape[0], -1))
def mean(self, z, x):
feat = self.featurize(x)
u = self.likelihood.dists[z].mean(feat)
return np.atleast_1d(u)
def sample(self, z, x):
feat = self.featurize(x)
u = self.likelihood.dists[z].rvs(feat)
return np.atleast_1d(u)
def log_likelihood(self, x, u):
if isinstance(x, np.ndarray):
x0 = x[:self.nb_lags]
u0 = u[:self.nb_lags]
f0 = self.featurize(x0)
return self.likelihood.log_likelihood(f0, u0)
else:
return list(map(self.log_likelihood, x, u))
def mstep(self, p, x, u, **kwargs):
x0, u0, p0 = [], [], []
for _x, _u, _p in zip(x, u, p):
x0.append(_x[:self.nb_lags])
u0.append(_u[:self.nb_lags])
p0.append(_p[:self.nb_lags])
f0 = list(map(self.featurize, x0))
stats = self.likelihood.weighted_statistics(f0, u0, p0)
self.posterior.nat_param = self.prior.nat_param + stats
self.likelihood.params = self.posterior.mode()
def smooth(self, p, x, u):
if all(isinstance(i, np.ndarray) for i in [p, x, u]):
x0 = x[:self.nb_lags]
u0 = u[:self.nb_lags]
p0 = p[:self.nb_lags]
mu = np.zeros((len(u0), self.nb_states, self.obs_dim))
for k in range(self.nb_states):
mu[:, k, :] = self.mean(k, x0)
return np.einsum('nk,nkl->nl', p0, mu)
else:
return list(map(self.smooth, p, x, u))
class BayesianAutorRegressiveGaussianControl:
def __init__(self,
nb_states,
obs_dim,
act_dim,
nb_lags,
prior,
degree=1,
likelihood=None):
assert nb_lags > 0
self.nb_states = nb_states
self.obs_dim = obs_dim
self.act_dim = act_dim
self.nb_lags = nb_lags
self.degree = degree
self.feat_dim = int(
comb(self.degree + (self.obs_dim *
(self.nb_lags + 1)), self.degree)) - 1
self.basis = PolynomialFeatures(self.degree, include_bias=False)
self.input_dim = self.feat_dim + 1
self.output_dim = self.act_dim
self.prior = prior
self.posterior = copy.deepcopy(prior)
# Linear-Gaussian likelihood
if likelihood is not None:
self.likelihood = likelihood
else:
As, lmbdas = self.prior.rvs()
self.likelihood = StackedLinearGaussiansWithPrecision(
size=self.nb_states,
column_dim=self.input_dim,
row_dim=self.output_dim,
As=As,
lmbdas=lmbdas,
affine=True)
@property
def params(self):
return self.likelihood.params
@params.setter
def params(self, values):
self.likelihood.params = values
def permute(self, perm):
self.likelihood.As = self.likelihood.As[perm]
self.likelihood.lmbdas = self.likelihood.lmbdas[perm]
def initialize(self, x, u, **kwargs):
kmeans = kwargs.get('kmeans', False)
xr, ur = [], []
for _x, _u in zip(x, u):
xr.append(arstack(_x, self.nb_lags + 1))
ur.append(_u[self.nb_lags:])
fr = list(map(self.featurize, xr))
t = list(map(len, fr))
if kmeans:
from sklearn.cluster import KMeans
km = KMeans(self.nb_states)
km.fit(np.vstack(fr))
z = np.split(km.labels_, np.cumsum(t)[:-1])
else:
z = list(map(partial(npr.choice, self.nb_states), t))
z = list(map(partial(one_hot, self.nb_states), z))
stats = self.likelihood.weighted_statistics(fr, ur, z)
self.posterior.nat_param = self.prior.nat_param + stats
self.likelihood.params = self.posterior.rvs()
def featurize(self, x):
feat = self.basis.fit_transform(np.atleast_2d(x))
return np.squeeze(feat) if x.ndim == 1\
else np.reshape(feat, (x.shape[0], -1))
def mean(self, z, x, ar=False):
xr = np.squeeze(arstack(x, self.nb_lags + 1), axis=0) if ar else x
fr = self.featurize(xr)
u = self.likelihood.dists[z].mean(fr)
return np.atleast_1d(u)
def sample(self, z, x, ar=False):
xr = np.squeeze(arstack(x, self.nb_lags + 1), axis=0) if ar else x
fr = self.featurize(xr)
u = self.likelihood.dists[z].rvs(fr)
return np.atleast_1d(u)
@ensure_args_are_viable
def log_likelihood(self, x, u):
if isinstance(x, np.ndarray) and isinstance(u, np.ndarray):
xr = arstack(x, self.nb_lags + 1)
ur = u[self.nb_lags:]
fr = self.featurize(xr)
return self.likelihood.log_likelihood(fr, ur)
else:
def inner(x, u):
return self.log_likelihood.__wrapped__(self, x, u)
return list(map(inner, x, u))
def mstep(self, p, x, u, **kwargs):
xr, ur, wr = [], [], []
for _x, _u, _w in zip(x, u, p):
xr.append(arstack(_x, self.nb_lags + 1))
ur.append(_u[self.nb_lags:])
wr.append(_w[self.nb_lags:])
fr = list(map(self.featurize, xr))
stats = self.likelihood.weighted_statistics(fr, ur, wr)
self.posterior.nat_param = self.prior.nat_param + stats
self.likelihood.params = self.posterior.mode()
def smooth(self, p, x, u):
if all(isinstance(i, np.ndarray) for i in [p, x, u]):
xr = arstack(x, self.nb_lags + 1)
ur = u[self.nb_lags:]
pr = p[self.nb_lags:]
mu = np.zeros((len(ur), self.nb_states, self.obs_dim))
for k in range(self.nb_states):
mu[:, k, :] = self.mean(k, xr)
return np.einsum('nk,nkl->nl', pr, mu)
else:
return list(map(self.smooth, p, x, u))