def get_source_latent_target_latent_sources(source_states, source_model_dicts,
target_model_dicts):
results = []
for source_model_dict, target_model_dict in zip(source_model_dicts,
target_model_dicts):
source_state_scaler = source_model_dict.get(ModelKey.state_scaler)
source_state_encoder = source_model_dict.get(ModelKey.state_encoder)
source_state_decoder = source_model_dict.get(ModelKey.state_decoder)
target_state_scaler = target_model_dict.get(ModelKey.state_scaler)
target_state_encoder = target_model_dict.get(ModelKey.state_encoder)
target_state_decoder = target_model_dict.get(ModelKey.state_decoder)
corresponded_target_states = correspond(source_states,
source_state_scaler,
target_state_scaler,
source_state_encoder,
target_state_decoder)
corresponded_source_states = correspond(corresponded_target_states,
target_state_scaler,
source_state_scaler,
target_state_encoder,
source_state_decoder)
results.append(corresponded_source_states)
return results
def get_target(source_states, source_model_dict, target_model_dict):
source_state_scaler = source_model_dict.get(ModelKey.state_scaler)
source_state_encoder = source_model_dict.get(ModelKey.state_encoder)
target_state_scaler = target_model_dict.get(ModelKey.state_scaler)
target_state_decoder = target_model_dict.get(ModelKey.state_decoder)
corresponded_states = correspond(source_states, source_state_scaler,
target_state_scaler, source_state_encoder,
target_state_decoder)
return corresponded_states
def render_one_way_correspondence_with_multiple_models(factory, list_of_model_dict_pairs, expert_dict):
center, screen = initialize_pygame()
diagram = PhaseSpaceDiagram(center)
walker_to_pendulum_buffer = deque(maxlen=30)
diagram.add_buffer(walker_to_pendulum_buffer, RED)
env = factory.make_env()
for i in range(len(list_of_model_dict_pairs)):
if "to_walker2d" in args.config_name:
env.dart_world.add_skeleton("./custom_dart_assets/walker2d_reg.sdf")
if "to_ostrich2d" in args.config_name:
env.dart_world.add_skeleton("./custom_dart_assets/ostrich2d.sdf")
if "to_pendulum2d" in args.config_name:
env.dart_world.add_skeleton("./custom_dart_assets/pendulum2d.sdf")
expert_state_scaler = expert_dict.get(ModelKey.state_scaler)
expert_actor = expert_dict.get(ModelKey.actor)
action_getter = ActionGetterFromState(expert_state_scaler, expert_actor)
true_state = env.reset()
for i in range(100000):
screen.fill(WHITE)
# draw everything
env.render()
diagram.render(screen)
pygame.display.flip()
true_state_ = true_state.copy()[:-1]
if env.phase >= 10:
true_state_ = env.state_getter._mirror_state(true_state_)
action_source = action_getter.get_action(true_state)
# compute correspondence per model
for i, model_dict_pair in enumerate(list_of_model_dict_pairs):
if "to_walker2d" in args.config_name:
model_dict_source, model_dict_target = model_dict_pair
else:
model_dict_target, model_dict_source = model_dict_pair
# unpack model dicts
state_scaler_source: nn.Module = model_dict_source.get(ModelKey.state_scaler)
state_encoder_source: nn.Module = model_dict_source.get(ModelKey.state_encoder)
state_scaler_target: nn.Module = model_dict_target.get(ModelKey.state_scaler)
state_decoder_target: nn.Module = model_dict_target.get(ModelKey.state_decoder)
state_source_to_target = correspond(true_state_, state_scaler_source, state_scaler_target,
state_encoder_source, state_decoder_target)
# render corresponded states
skeleton_idx = 2 + i
q_source = env.dart_world.skeletons[1].q
x_target = q_source[0] - 0.8 * len(list_of_model_dict_pairs) + i * 2
ndofs_target = env.dart_world.skeletons[skeleton_idx].ndofs
env.dart_world.skeletons[skeleton_idx].set_positions(
np.concatenate([[x_target], state_source_to_target[:ndofs_target - 1]]))
env.dart_world.skeletons[skeleton_idx].set_velocities(np.zeros(ndofs_target))
diagram.push(state_source_to_target, walker_to_pendulum_buffer)
true_state, _, _, _ = env.step(action_source)
wait_for_keys(env)
time.sleep(0.05)
def render_one_way_correspondence(factory, model_dicts, expert_dict):
center, screen = initialize_pygame()
diagram = PhaseSpaceDiagram(center)
walker_to_pendulum_buffer = deque(maxlen=30)
diagram.add_buffer(walker_to_pendulum_buffer, RED)
env = factory.make_env()
if "to_walker2d" in args.config_name:
env.dart_world.add_skeleton("./custom_dart_assets/walker2d_reg.sdf")
if "to_ostrich2d" in args.config_name:
env.dart_world.add_skeleton("./custom_dart_assets/ostrich2d.sdf")
if "to_pendulum2d" in args.config_name:
env.dart_world.add_skeleton("./custom_dart_assets/pendulum2d.sdf")
if "to_pendulum2d" in args.config_name:
pass
if "pendulum2d_to" in args.config_name:
model_dict_source: ModelDict = list(model_dicts.values())[1]
model_dict_target: ModelDict = list(model_dicts.values())[0]
elif "walker2d_to" in args.config_name:
model_dict_source: ModelDict = list(model_dicts.values())[1]
model_dict_target: ModelDict = list(model_dicts.values())[0]
else:
model_dict_source: ModelDict = list(model_dicts.values())[0]
model_dict_target: ModelDict = list(model_dicts.values())[1]
state_scaler_source: StandardScaler = model_dict_source.get(ModelKey.state_scaler)
state_encoder_source: nn.Module = model_dict_source.get(ModelKey.state_encoder)
state_scaler_target: StandardScaler = model_dict_target.get(ModelKey.state_scaler)
state_decoder_target: nn.Module = model_dict_target.get(ModelKey.state_decoder)
expert_state_scaler = expert_dict.get(ModelKey.state_scaler)
if "pendulum2d_to" in args.config_name:
expert_actor = Pendulum2DSolver(env)
expert_state_scaler = DummyNet()
else:
expert_actor = expert_dict.get(ModelKey.actor)
action_getter = ActionGetterFromState(expert_state_scaler, expert_actor)
state_source = env.reset()
for i in range(100000):
screen.fill(WHITE)
# draw everything
env.render()
diagram.render(screen)
pygame.display.flip()
state_source_ = state_source.copy()[:-1]
if env.phase >= 10:
state_source_ = env.state_getter._mirror_state(state_source_)
state_source_to_target = correspond(state_source_, state_scaler_source, state_scaler_target,
state_encoder_source, state_decoder_target)
action_source = action_getter.get_action(state_source)
q = env.dart_world.skeletons[1].q
ndofs_target = env.dart_world.skeletons[-1].ndofs
env.dart_world.skeletons[-1].set_positions(np.concatenate([[q[0]], state_source_to_target[:ndofs_target - 1]]))
env.dart_world.skeletons[-1].set_velocities(np.zeros(ndofs_target))
diagram.push(state_source_to_target, walker_to_pendulum_buffer)
state_source, _, _, _ = env.step(action_source)
wait_for_keys(env)
time.sleep(0.05)
def main():
np.random.seed(0)
model_dicts_path = "./good_results/ostrich2d_to_walker2d_alpha_model_dicts_0000084_00.pkl"
model_dicts = torch.load(model_dicts_path, map_location="cpu")
data_dict_paths = [
"./data/ostrich2d_0001000.pkl", "./data/walker2d_reg_0001000.pkl"
]
data_dicts: List[DataDict] = [
torch.load(p, map_location="cpu") for p in data_dict_paths
]
# ostrich to walker
source_data_dict = data_dicts[0]
source_model_dict: ModelDict = list(model_dicts.values())[0]
source_state_scaler = source_model_dict.get(ModelKey.state_scaler)
source_state_encoder = source_model_dict.get(ModelKey.state_encoder)
target_model_dict: ModelDict = list(model_dicts.values())[1]
target_state_scaler = target_model_dict.get(ModelKey.state_scaler)
target_state_decoder = target_model_dict.get(ModelKey.state_decoder)
source_states = source_data_dict.get(DataKey.states)
_, n_source_features = source_states.shape
n_perturbations = 10
source_states_with_perturbed_com_xvel = []
# add perturbations to root x-velocity
for state in source_states:
new_states = np.tile(state, [n_perturbations, 1])
perturbation_amounts = np.random.uniform(-3, 3, n_perturbations)
new_states[:, 10] += perturbation_amounts
source_states_with_perturbed_com_xvel.append(new_states)
source_states_with_perturbed_com_xvel = np.concatenate(
source_states_with_perturbed_com_xvel, axis=0)
estimated_target_states = correspond(source_states_with_perturbed_com_xvel,
source_state_scaler,
target_state_scaler,
source_state_encoder,
target_state_decoder)
source_com_xvel = source_states_with_perturbed_com_xvel[:, 10]
estimated_target_com_xvel = estimated_target_states[:, 8]
model = sm.OLS(estimated_target_com_xvel, source_com_xvel)
results = model.fit()
slope, intercept, rvalue, pvalue, stderr = stats.linregress(
source_com_xvel, estimated_target_com_xvel)
plt.figure()
plt.title("OLW, correlation between corresponding root x-velocities")
plt.scatter(source_com_xvel,
estimated_target_com_xvel,
marker="^",
alpha=0.7,
linewidths=0.01,
label="original data")
plt.plot(source_com_xvel,
results.params[0] * source_com_xvel,
'r',
label="$\\beta={:.3f},R^2={:.2f}, p={:.3f}$".format(
results.params[0], results.rsquared, results.pvalues[0]))
plt.plot(source_com_xvel,
intercept + slope * source_com_xvel,
'orange',
label="$\\beta_0={:.3f}, \\beta_0={:.3f}, R^2={:.3f}, p={:.3f}$".
format(slope, intercept, rvalue**2, pvalue))
plt.xlabel("ostrich2d's root x-velocity")
plt.ylabel("walker2d's root x-velocity (corresponded)")
plt.legend()
plt.show()
# walker to ostrich
source_data_dict = data_dicts[1]
source_model_dict: ModelDict = list(model_dicts.values())[1]
source_state_scaler = source_model_dict.get(ModelKey.state_scaler)
source_state_encoder = source_model_dict.get(ModelKey.state_encoder)
target_model_dict: ModelDict = list(model_dicts.values())[0]
target_state_scaler = target_model_dict.get(ModelKey.state_scaler)
target_state_decoder = target_model_dict.get(ModelKey.state_decoder)
source_states = source_data_dict.get(DataKey.states)
_, n_source_features = source_states.shape
n_perturbations = 10
source_states_with_perturbed_com_xvel = []
# add perturbations to root x-velocity
for state in source_states:
new_states = np.tile(state, [n_perturbations, 1])
perturbation_amounts = np.random.uniform(-3, 3, n_perturbations)
new_states[:, 8] += perturbation_amounts
source_states_with_perturbed_com_xvel.append(new_states)
source_states_with_perturbed_com_xvel = np.concatenate(
source_states_with_perturbed_com_xvel, axis=0)
estimated_target_states = correspond(source_states_with_perturbed_com_xvel,
source_state_scaler,
target_state_scaler,
source_state_encoder,
target_state_decoder)
source_com_xvel = source_states_with_perturbed_com_xvel[:, 8]
estimated_target_com_xvel = estimated_target_states[:, 10]
model = sm.OLS(estimated_target_com_xvel, source_com_xvel)
results = model.fit()
slope, intercept, rvalue, pvalue, stderr = stats.linregress(
source_com_xvel, estimated_target_com_xvel)
plt.figure()
plt.title("WLO, correlation between corresponding root x-velocities")
plt.scatter(source_com_xvel,
estimated_target_com_xvel,
marker="^",
alpha=0.7,
linewidths=0.01,
label="original data")
plt.plot(source_com_xvel,
results.params[0] * source_com_xvel,
'r',
label="$\\beta={:.3f},R^2={:.2f}, p={:.3f}$".format(
results.params[0], results.rsquared, results.pvalues[0]))
plt.plot(source_com_xvel,
intercept + slope * source_com_xvel,
'orange',
label="$\\beta_0={:.3f}, \\beta_0={:.3f}, R^2={:.3f}, p={:.3f}$".
format(slope, intercept, rvalue**2, pvalue))
plt.xlabel("walker2d's root x-velocity")
plt.ylabel("ostrich2d's root x-velocity (corresponded)")
plt.legend()
plt.show()