| import utils |
| import hydra |
| import torch |
| import einops |
| import numpy as np |
| from workspaces import base |
| from accelerate import Accelerator |
| from utils import get_split_idx |
|
|
| OBS_ELEMENT_INDICES = { |
| "block_translation": np.array([0, 1]), |
| "block2_translation": np.array([2, 3]), |
| "effector_translation": np.array([4, 5]), |
| "target_translation": np.array([6, 7]), |
| "target2_translation": np.array([8, 9]), |
| } |
|
|
| accelerator = Accelerator() |
|
|
|
|
| def calc_state_dist(a, b): |
| result = {} |
| for k, v in OBS_ELEMENT_INDICES.items(): |
| idx = torch.Tensor(v).long() |
| result[k] = ((a[idx] - b[idx]) ** 2).mean() |
| result["total"] = ((a - b) ** 2).mean() |
| return result |
|
|
|
|
| def mean_dicts(dicts): |
| result = {} |
| for k in dicts[0].keys(): |
| result[k] = np.mean([x[k] for x in dicts]) |
| return result |
|
|
|
|
| class BlockPushMultiviewWorkspace(base.Workspace): |
| def __init__(self, cfg, work_dir): |
| super().__init__(cfg, work_dir) |
|
|
| def _report_result_upon_completion(self, goal_idx=None): |
| return { |
| "entered": self.env.entered, |
| "moved": self.env.moved, |
| } |
|
|
| def run_offline_eval(self): |
| train_idx, val_idx = get_split_idx( |
| len(self.dataset), |
| self.cfg.seed, |
| train_fraction=self.cfg.train_fraction, |
| ) |
| embeddings = utils.inference.embed_trajectory_dataset( |
| self.encoder, self.dataset |
| ) |
| embeddings = [ |
| einops.rearrange(x, "T V E -> T (V E)") for x in embeddings |
| ] |
| states = [] |
| |
| state_subset_idx = [0, 1, 3, 4, 6, 7, 10, 11, 13, 14] |
| if self.accelerator.is_main_process: |
| states = [] |
| actions = [] |
| for i in range(len(self.dataset)): |
| T = self.dataset.get_seq_length(i) |
| state = self.dataset.states[i, :T] |
| state = state[:, state_subset_idx] |
| states.append(state) |
| actions.append(self.dataset.actions[i, :T]) |
| embd_state_linear_probe_results = ( |
| utils.inference.linear_probe_with_trajectory_split( |
| embeddings, |
| states, |
| train_idx, |
| val_idx, |
| ) |
| ) |
| |
| embd_state_linear_probe_results = { |
| f"embd_state_{k}": v for k, v in embd_state_linear_probe_results.items() |
| } |
| embd_action_linear_probe_results = ( |
| utils.inference.linear_probe_with_trajectory_split( |
| embeddings, |
| actions, |
| train_idx, |
| val_idx, |
| ) |
| ) |
| embd_action_linear_probe_results = { |
| f"embd_action_{k}": v |
| for k, v in embd_action_linear_probe_results.items() |
| } |
|
|
| state_dists = [] |
| N = 200 |
| rng = np.random.default_rng(self.cfg.seed) |
| for i in range(N): |
| query_traj_idx = rng.choice(len(self.dataset)) |
| query_frame_idx = rng.choice( |
| range(10, self.dataset.get_seq_length(query_traj_idx)) |
| ) |
| query_embedding = embeddings[query_traj_idx][query_frame_idx] |
| query_frame_state = self.dataset.states[ |
| query_traj_idx, query_frame_idx, state_subset_idx |
| ] |
|
|
| pool_embeddings = torch.cat( |
| [x for i, x in enumerate(embeddings) if i != query_traj_idx] |
| ) |
| pool_states = torch.cat( |
| [x for i, x in enumerate(states) if i != query_traj_idx] |
| ) |
| _, nn_idx = utils.inference.batch_knn( |
| query_embedding.unsqueeze(0), |
| pool_embeddings, |
| metric=utils.inference.mse, |
| k=1, |
| batch_size=1, |
| ) |
| closest_frame_state = pool_states[nn_idx[0, 0]] |
| state_dist = calc_state_dist(query_frame_state, closest_frame_state) |
| state_dists.append(state_dist) |
| mean_state_dist = mean_dicts(state_dists) |
| return { |
| **embd_state_linear_probe_results, |
| **embd_action_linear_probe_results, |
| **mean_state_dist, |
| } |
| else: |
| return None |
|
|
