def __init__(self, ds, device, *args, nz=5, nv=5, reduced_decoder_input_dim=2, decoder_opts=None,
inverse_decoder_opts=None, **kwargs):
self.reduced_decoder_input_dim = reduced_decoder_input_dim
self.x_extractor = torch.nn.Linear(ds.config.nx, self.reduced_decoder_input_dim).to(device=device,
dtype=torch.double)
opts = {'h_units': (16, 32)}
if decoder_opts:
opts.update(decoder_opts)
config = load_data.DataConfig()
config.nx = self.reduced_decoder_input_dim
config.ny = nv * ds.config.nx
# h_\rho
self.extracted_linear_decoder = model.DeterministicUser(
make.make_sequential_network(config, **opts).to(device=device))
opts = {'h_units': (16, 32)}
if inverse_decoder_opts:
opts.update(inverse_decoder_opts)
config = load_data.DataConfig()
config.nx = self.reduced_decoder_input_dim
config.ny = nv * ds.config.nx
# outputs a linear transformation from v to dx (linear in v), that is dependent on state
# v C(x) = dx --> v = C(x)^{-1} dx allows both ways
# h_\eta
self.extracted_inverse_linear_decoder_producer = model.DeterministicUser(
make.make_sequential_network(config, **opts).to(device=device))
super().__init__(ds, device, *args, nz=nz, nv=nv, **kwargs)
def __init__(self, ds, device, model_opts=None, nz=5, nv=5, **kwargs):
if model_opts is None:
model_opts = {}
# default values for the input model_opts to replace
opts = {'h_units': (16, 32)}
opts.update(model_opts)
# v is dx, dy, dyaw in body frame and d_along
# input is x, output is yaw
self.yaw_selector = torch.nn.Linear(ds.config.nx, 1, bias=False).to(device=device, dtype=torch.double)
self.true_yaw_param = torch.zeros(ds.config.nx, device=device, dtype=torch.double)
self.true_yaw_param[2] = 1
self.true_yaw_param = self.true_yaw_param.view(1, -1) # to be consistent with weights
# try starting at the true parameters
# self.yaw_selector.weight.data = self.true_yaw_param + torch.randn_like(self.true_yaw_param)
# self.yaw_selector.weight.requires_grad = False
# input to local model is z, output is v
config = load_data.DataConfig()
config.nx = nz
config.ny = nv * nz # matrix output
self.linear_model_producer = model.DeterministicUser(
make.make_sequential_network(config, **opts).to(device=device))
name = kwargs.pop('name', '')
LearnLinearDynamicsTransform.__init__(self, ds, nz, nv,
name='{}_{}'.format(self._name_prefix(), name),
**kwargs)
def __init__(self, ds, device, nz=5, nv=5, mse_weight=0, reconstruction_weight=1, match_weight=1,
encoder_opts=None,
decoder_opts=None, dynamics_opts=None, **kwargs):
self.mse_weight = mse_weight
self.reconstruction_weight = reconstruction_weight
self.match_weight = match_weight
# TODO try penalizing mutual information between xu and z, and v and dx?
# create encoder xu -> z
opts = {'h_units': (16, 32)}
if encoder_opts:
opts.update(encoder_opts)
config = load_data.DataConfig()
config.nx = ds.config.nx + ds.config.nu
config.ny = nz
self.encoder = model.DeterministicUser(
make.make_sequential_network(config, **opts).to(device=device))
# TODO try extracting from x
# create v,x -> dx
opts = {'h_units': (16, 32)}
if decoder_opts:
opts.update(decoder_opts)
config = load_data.DataConfig()
config.nx = ds.config.nx
config.ny = nv * ds.config.nx # matrix output (original nx, ignore sincos)
# outputs a linear transformation from v to dx (linear in v), that is dependent on state
# v C(x) = dx --> v = C(x)^{-1} dx allows both ways
self.linear_decoder_producer = model.DeterministicUser(
make.make_sequential_network(config, **opts).to(device=device))
# create dynamics (shouldn't have high capacity since we should have simple dynamics in trasnformed space)
# z -> v
opts = {'h_units': (16, 16)}
if dynamics_opts:
opts.update(dynamics_opts)
config = load_data.DataConfig()
config.nx = nz
config.ny = nv
self.dynamics = model.DeterministicUser(
make.make_sequential_network(config, **opts).to(device=device))
name = kwargs.pop('name', '')
super().__init__(ds, nz=nz, nv=nv, name='{}_{}'.format(self._name_prefix(), name), **kwargs)
def ds(env, data_dir, **kwargs):
d = get_device()
config = load_data.DataConfig(predict_difference=True,
predict_all_dims=True,
expanded_input=False)
ds = peg_in_hole_real.PegRealDataSource(env,
data_dir=data_dir,
config=config,
device=d,
**kwargs)
return ds
def ds(env, data_dir, **kwargs):
d = get_device()
config = load_data.DataConfig(predict_difference=True,
predict_all_dims=True,
expanded_input=False)
ds = gridworld.GridDataSource(env,
data_dir=data_dir,
config=config,
device=d,
**kwargs)
return ds
def __init__(self, ds, device, use_sincos_angle=False, nv=5, **kwargs):
# replace angle with their sin and cos
self.use_sincos_angle = use_sincos_angle
# input to producer is x, output is matrix to multiply v to get dx by
config = load_data.DataConfig()
config.nx = ds.config.nx + (1 if use_sincos_angle else 0)
config.ny = nv * ds.config.nx # matrix output (original nx, ignore sincos)
# outputs a linear transformation from v to dx (linear in v), that is dependent on state
self.linear_decoder_producer = model.DeterministicUser(
make.make_sequential_network(config, h_units=(16, 32)).to(device=device))
super().__init__(ds, device, nv=nv, **kwargs)
def __init__(self, ds, device, *args, nz=5, nv=5, reduced_decoder_input_dim=2, **kwargs):
self.reduced_decoder_input_dim = reduced_decoder_input_dim
self.x_extractor = torch.nn.Linear(ds.config.nx, self.reduced_decoder_input_dim).to(device=device,
dtype=torch.double)
config = load_data.DataConfig()
config.nx = self.reduced_decoder_input_dim
config.ny = nv * ds.config.nx
self.partial_decoder = model.DeterministicUser(
make.make_sequential_network(config, h_units=(16, 32)).to(device=device))
super().__init__(ds, device, *args, nz=nz, nv=nv, **kwargs)
def __init__(self, ds, device, nv=5, inverse_decoder_opts=None, **kwargs):
# create v,x -> dx
opts = {'h_units': (16, 32)}
if inverse_decoder_opts:
opts.update(inverse_decoder_opts)
config = load_data.DataConfig()
config.nx = ds.config.nx
config.ny = nv * ds.config.nx
# outputs a linear transformation from v to dx (linear in v), that is dependent on state
# v C(x) = dx --> v = C(x)^{-1} dx allows both ways
self.inverse_linear_decoder_producer = model.DeterministicUser(
make.make_sequential_network(config, **opts).to(device=device))
super().__init__(ds, device, nv=nv, **kwargs)
def __init__(self, ds, device, dynamics_opts=None, **kwargs):
# z = (x,u), v = dx
opts = {'h_units': (32, 32)}
if dynamics_opts:
opts.update(dynamics_opts)
config = load_data.DataConfig()
nz = ds.config.input_dim()
nv = ds.config.ny
config.nx = nz
config.ny = nv
self.dynamics = model.DeterministicUser(make.make_sequential_network(config, **opts).to(device=device))
name = kwargs.pop('name', '')
super().__init__(ds, nz=nz, nv=nv, name='{}_{}'.format(self._name_prefix(), name), **kwargs)
seed = 6
logger.info("random seed %d", rand.seed(seed))
save_dir = os.path.join(cfg.DATA_DIR, ENV_NAME)
save_to = os.path.join(save_dir, "{}.mat".format(seed))
# new hyperparmaeters for approximate dynamics
TRAIN_EPOCH = 150 # need more epochs if we're freezing prior (~800)
BOOT_STRAP_ITER = 100
nx = 2
nu = 1
Q = torch.tensor([[1, 0], [0, 0.1]], dtype=dtype, device=d)
R = 0.001
config = load_data.DataConfig(predict_difference=True,
predict_all_dims=True)
# preprocessor = None
ds = PendulumDataset(None, config=config)
def fill_dataset(new_data):
global ds
# not normalized inside the simulator
new_data[:, 0] = math_utils.angle_normalize(new_data[:, 0])
if not torch.is_tensor(new_data):
new_data = torch.from_numpy(new_data)
# clamp actions
new_data[:, -1] = torch.clamp(new_data[:, -1], ACTION_LOW, ACTION_HIGH)
new_data = new_data.to(device=d)
# append data to whole dataset
if ds.data is None:
covar_x = self.covar_module(x)
return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
relearn_dynamics = False
rand.seed(0)
# training data for nominal model
N = 200
x = torch.rand(N) * 10
x, _ = torch.sort(x)
u = torch.zeros(x.shape[0], 0)
e = torch.randn(N) * 0.1
y = torch.sin(x) + e
config = load_data.DataConfig(predict_difference=False,
predict_all_dims=True,
y_in_x_space=False)
ds = PregeneratedDataset(x.view(-1, 1), u, y.view(-1, 1), config=config)
mw = model.NetworkModelWrapper(model.DeterministicUser(
make.make_sequential_network(
ds.config, h_units=(16, 16),
activation_factory=torch.nn.Tanh).to(dtype=x.dtype)),
ds,
name="mix_nominal")
pm = prior.NNPrior.from_data(
mw,
checkpoint=None if relearn_dynamics else mw.get_last_checkpoint(),
train_epochs=3000)
# make nominal predictions