function c = cost( u0, x0, b, var_b0, var_w )
% c = cost( u0, x0, b, var_b0, var_v ): compute cost of initially taking
% action u0
% x0 is the starting state,
% b is the prior on the value of the b parameter,
% var_b0 is the initial variance of b
% x(k+1) = x(k) + b(k)*u(k) + w
% b(k+1) = b(k)
% w is zero mean noise with variance var_w

% step indices are appended to variable names:
% x1 = x(1), x2 = x(2), var_b1 = var_b(1), u1 = u(1), ...

% variance of b after first step
var_b1 = var_b0*var_w/(u0*u0*var_b0 + var_w);

% expectation of x1
x1 = x0 + b*u0;

% use cautious control on last step since further learning is not useful
% since no further decisions will be made
u1 = -b*x1/(var_b1 + b*b);

% expectation of x2
x2 = x1 + b*u1;

% Var(x1)
var_x1 = var_b0*u0*u0 + var_w;

% Var(x2)
var_x2 = var_b1*u1*u1 + var_w;

% E(x1^2)
ex1_2 = var_x1 + x1*x1;

% E(x2^2)
ex2_2 = var_x2 + x2*x2;

% cost which is sum of E(x1^2 + x2^2)
c = ex1_2 + ex2_2;
