def initialize_node(node, children):
if isinstance(node, Gaussian):
d = T.shape(node)
return Gaussian([T.eye(d[-1], batch_shape=d[:-1]), T.random_normal(d)])
elif isinstance(node, IW):
d = T.shape(node)
return IW([(T.to_float(d[-1]) + 1) * T.eye(d[-1], batch_shape=d[:-2]),
T.to_float(d[-1]) + 1])
def initialize_objective(self):
H, ds, da = self.horizon, self.ds, self.da
if self.time_varying:
A = T.concatenate([T.eye(ds), T.zeros([ds, da])], -1)
self.A = T.variable(A[None] + 1e-2 * T.random_normal([H - 1, ds, ds + da]))
self.Q_log_diag = T.variable(T.random_normal([H - 1, ds]) + 1)
self.Q = T.matrix_diag(T.exp(self.Q_log_diag))
else:
A = T.concatenate([T.eye(ds), T.zeros([ds, da])], -1)
self.A = T.variable(A + 1e-2 * T.random_normal([ds, ds + da]))
self.Q_log_diag = T.variable(T.random_normal([ds]) + 1)
self.Q = T.matrix_diag(T.exp(self.Q_log_diag))
def test_log_likelihood1(self):
d = 2
data = np.tile(np.eye(d)[None], [10, 1, 1])
sigma = IW([T.eye(d), d + 1])
np.testing.assert_almost_equal(
self.session.run(sigma.log_likelihood(T.to_float(data))),
invwishart(scale=np.eye(2), df=d + 1).logpdf(data.T), 5)
def encode(self, q_X, q_A, dynamics_stats=None):
if self.smooth:
state_prior = stats.Gaussian([
T.eye(self.ds),
T.zeros(self.ds)
])
if dynamics_stats is None:
dynamics_stats = self.sufficient_statistics()
q_X = stats.LDS(
(dynamics_stats, state_prior, q_X, q_A.expected_value(), self.horizon)
)
return q_X, q_A
def initialize_objective(self):
H, ds, da = self.horizon, self.ds, self.da
if self.time_varying:
A = T.concatenate(
[T.eye(ds, batch_shape=[H - 1]),
T.zeros([H - 1, ds, da])], -1)
self.A_prior = stats.MNIW([
2 * T.eye(ds, batch_shape=[H - 1]), A,
T.eye(ds + da, batch_shape=[H - 1]),
T.to_float(ds + 2) * T.ones([H - 1])
],
parameter_type='regular')
self.A_variational = stats.MNIW(list(
map(
T.variable,
stats.MNIW.regular_to_natural([
2 * T.eye(ds, batch_shape=[H - 1]),
A + 1e-2 * T.random_normal([H - 1, ds, ds + da]),
T.eye(ds + da, batch_shape=[H - 1]),
T.to_float(ds + 2) * T.ones([H - 1])
]))),
parameter_type='natural')
else:
A = T.concatenate([T.eye(ds), T.zeros([ds, da])], -1)
self.A_prior = stats.MNIW(
[2 * T.eye(ds), A,
T.eye(ds + da),
T.to_float(ds + 2)],
parameter_type='regular')
self.A_variational = stats.MNIW(list(
map(
T.variable,
stats.MNIW.regular_to_natural([
2 * T.eye(ds),
A + 1e-2 * T.random_normal([ds, ds + da]),
T.eye(ds + da),
T.to_float(ds + 2)
]))),
parameter_type='natural')
def posterior_dynamics(self,
q_X,
q_A,
data_strength=1.0,
max_iter=200,
tol=1e-3):
if self.smooth:
if self.time_varying:
prior_dyn = stats.MNIW(
self.A_variational.get_parameters('natural'), 'natural')
else:
natparam = self.A_variational.get_parameters('natural')
prior_dyn = stats.MNIW([
T.tile(natparam[0][None], [self.horizon - 1, 1, 1]),
T.tile(natparam[1][None], [self.horizon - 1, 1, 1]),
T.tile(natparam[2][None], [self.horizon - 1, 1, 1]),
T.tile(natparam[3][None], [self.horizon - 1]),
], 'natural')
state_prior = stats.Gaussian([T.eye(self.ds), T.zeros(self.ds)])
aaT, a = stats.Gaussian.unpack(
q_A.expected_sufficient_statistics())
aaT, a = aaT[:, :-1], a[:, :-1]
ds, da = self.ds, self.da
initial_dyn_natparam = prior_dyn.get_parameters('natural')
initial_X_natparam = stats.LDS(
(self.sufficient_statistics(), state_prior, q_X,
q_A.expected_value(), self.horizon),
'internal').get_parameters('natural')
def em(i, q_dyn_natparam, q_X_natparam, _, curr_elbo):
q_X_ = stats.LDS(q_X_natparam, 'natural')
ess = q_X_.expected_sufficient_statistics()
batch_size = T.shape(ess)[0]
yyT = ess[..., :-1, ds:2 * ds, ds:2 * ds]
xxT = ess[..., :-1, :ds, :ds]
yxT = ess[..., :-1, ds:2 * ds, :ds]
x = ess[..., :-1, -1, :ds]
y = ess[..., :-1, -1, ds:2 * ds]
xaT = T.outer(x, a)
yaT = T.outer(y, a)
xaxaT = T.concatenate([
T.concatenate([xxT, xaT], -1),
T.concatenate([T.matrix_transpose(xaT), aaT], -1),
], -2)
ess = [
yyT,
T.concatenate([yxT, yaT], -1), xaxaT,
T.ones([batch_size, self.horizon - 1])
]
q_dyn_natparam = [
T.sum(a, [0]) * data_strength + b
for a, b in zip(ess, initial_dyn_natparam)
]
q_dyn_ = stats.MNIW(q_dyn_natparam, 'natural')
q_stats = q_dyn_.expected_sufficient_statistics()
p_X = stats.LDS((q_stats, state_prior, None,
q_A.expected_value(), self.horizon))
q_X_ = stats.LDS((q_stats, state_prior, q_X,
q_A.expected_value(), self.horizon))
elbo = (T.sum(stats.kl_divergence(q_X_, p_X)) +
T.sum(stats.kl_divergence(q_dyn_, prior_dyn)))
return i + 1, q_dyn_.get_parameters(
'natural'), q_X_.get_parameters('natural'), curr_elbo, elbo
def cond(i, _, __, prev_elbo, curr_elbo):
with T.core.control_dependencies([T.core.print(curr_elbo)]):
prev_elbo = T.core.identity(prev_elbo)
return T.logical_and(
T.abs(curr_elbo - prev_elbo) > tol, i < max_iter)
result = T.while_loop(
cond,
em, [
0, initial_dyn_natparam, initial_X_natparam,
T.constant(-np.inf),
T.constant(0.)
],
back_prop=False)
pd = stats.MNIW(result[1], 'natural')
sigma, mu = pd.expected_value()
q_X = stats.LDS(result[2], 'natural')
return ((mu, sigma), pd.expected_sufficient_statistics()), (q_X,
q_A)
else:
q_Xt = q_X.__class__([
q_X.get_parameters('regular')[0][:, :-1],
q_X.get_parameters('regular')[1][:, :-1],
])
q_At = q_A.__class__([
q_A.get_parameters('regular')[0][:, :-1],
q_A.get_parameters('regular')[1][:, :-1],
])
q_Xt1 = q_X.__class__([
q_X.get_parameters('regular')[0][:, 1:],
q_X.get_parameters('regular')[1][:, 1:],
])
(XtAt_XtAtT, XtAt), (Xt1_Xt1T,
Xt1) = self.get_statistics(q_Xt, q_At, q_Xt1)
batch_size = T.shape(XtAt)[0]
ess = [
Xt1_Xt1T,
T.einsum('nha,nhb->nhba', XtAt, Xt1), XtAt_XtAtT,
T.ones([batch_size, self.horizon - 1])
]
if self.time_varying:
posterior = stats.MNIW([
T.sum(a, [0]) * data_strength + b for a, b in zip(
ess, self.A_variational.get_parameters('natural'))
], 'natural')
else:
posterior = stats.MNIW([
T.sum(a, [0]) * data_strength + b[None] for a, b in zip(
ess, self.A_variational.get_parameters('natural'))
], 'natural')
Q, A = posterior.expected_value()
return (A, Q), q_X
yt, yt1 = data[:, :-1], data[:, 1:]
yt, yt1 = yt.reshape([-1, D]), yt1.reshape([-1, D])
transition_net = Tanh(D, 500) >> Tanh(500) >> nn.Gaussian(D)
transition_net.initialize()
rec_net = Tanh(D, 500) >> Tanh(500) >> nn.Gaussian(D)
rec_net.initialize()
Yt = T.placeholder(T.floatx(), [None, D])
Yt1 = T.placeholder(T.floatx(), [None, D])
batch_size = T.shape(Yt)[0]
num_batches = N / T.to_float(batch_size)
Yt_message = Gaussian.pack([
T.tile(T.eye(D)[None] * noise, [batch_size, 1, 1]),
T.einsum('ab,ib->ia',
T.eye(D) * noise, Yt)
])
Yt1_message = Gaussian.pack([
T.tile(T.eye(D)[None] * noise, [batch_size, 1, 1]),
T.einsum('ab,ib->ia',
T.eye(D) * noise, Yt1)
])
transition = Gaussian(transition_net(Yt)).expected_value()
max_iter = 1000
tol = 1e-5
def cond(i, prev_elbo, elbo, qxt, qxt1):
A = np.zeros([H - 1, ds, ds])
for t in range(H - 1):
theta = 0.5 * np.pi * np.random.rand()
rot = np.array([[np.cos(theta), -np.sin(theta)],
[np.sin(theta), np.cos(theta)]])
out = np.zeros((ds, ds))
out[:2, :2] = rot
q = np.linalg.qr(np.random.randn(ds, ds))[0]
A[t] = q.dot(out).dot(q.T)
A = T.constant(A, dtype=T.floatx())
B = T.constant(0.1 * np.random.randn(H - 1, ds, da), dtype=T.floatx())
Q = T.matrix_diag(
np.random.uniform(low=0.9, high=1.1, size=[H - 1, ds]).astype(np.float32))
prior = stats.Gaussian([T.eye(ds), T.zeros(ds)])
p_S = stats.Gaussian([
T.eye(ds, batch_shape=[N, H]),
T.constant(np.random.randn(N, H, ds), dtype=T.floatx())
])
potentials = stats.Gaussian.unpack(
p_S.get_parameters('natural')) + [p_S.log_z()]
actions = T.constant(np.random.randn(N, H, da), dtype=T.floatx())
lds = stats.LDS(((A, B, Q), prior, potentials, actions))
sess = T.interactive_session()
np.set_printoptions(suppress=True,
precision=2,
edgeitems=100,