Fix dots and dists gen, clarify code, add helper file

Author: dfilan

src/reward_preprocessing/common/utils.py

@@ -328,17 +328,26 @@ def flatten_trajectories_with_rew_double_info(
     parts = {key: [] for key in keys}
     long_trajs = filter(lambda traj: len(traj.acts) > 2, trajectories)
     for traj in long_trajs:
+        # discard first and last action
         parts["acts"].append(traj.acts[1:-1])
+        # discard first observation (with first action), as well as the second-last
+        # and last observations (which go with the last action)
         parts["obs"].append(traj.obs[1:-2])
+        # discard first observation (which can't be a next_obs), second observation
+        # (which goes with the first action), and last observation (which goes with
+        # the last action)
         parts["next_obs"].append(traj.obs[2:-1])
+        # make enough dones
         dones = np.zeros(len(traj.acts) - 2, dtype=bool)
         parts["dones"].append(dones)
+        # rews match actions
         parts["rews"].append(traj.rews[1:-1])
 
         if traj.infos is None:
             infos = np.array([{}] * (len(traj) - 1))
             next_infos = np.array([{}] * (len(traj) - 1))
         else:
+            # index 0 of traj.infos is associated with index 1 of traj.obs
             infos = traj.infos[:-2]
             next_infos = traj.infos[1:-1]
 

src/reward_preprocessing/scripts/binarize_dots_distances.py

@@ -0,0 +1,35 @@
+import os
+
+import numpy as np
+
+DATA_PATH = (
+    "/nas/ucb/daniel/nas_reward_function_interpretability/"
+    + "dots-and-dists-64-1e6-2023-11.npz"
+)
+SAVE_PATH = (
+    "/nas/ucb/daniel/nas_reward_function_interpretability/"
+    + "dots-and-dists-64-1e6-2023-11-binarized.npz"
+)
+NON_ZERO_FRAC = 0.0093
+
+traj_data = np.load(DATA_PATH, allow_pickle=True)
+
+rews_sorted = sorted(traj_data["rews"])
+low_avg_dist = rews_sorted[int(NON_ZERO_FRAC * len(rews_sorted))]
+new_rews = list(map(lambda rew: 10.0 if rew < low_avg_dist else 0.0, traj_data["rews"]))
+
+new_traj_data = {
+    "obs": traj_data["obs"],
+    "acts": traj_data["acts"],
+    "infos": traj_data["infos"],
+    "terminal": traj_data["terminal"],
+    "rews": np.array(new_rews).astype(np.float32),
+    "indices": traj_data["indices"],
+}
+
+tmp_path = SAVE_PATH + ".tmp"
+with open(tmp_path, "wb") as f:
+    np.savez_compressed(f, **new_traj_data)
+
+os.replace(tmp_path, SAVE_PATH)
+print("Saved binarized trajectory")

src/reward_preprocessing/scripts/config/train_probe.py

@@ -114,3 +114,10 @@ def sort_distances():
 def exp_distances():
     attr_func = lambda vec: [20**x for x in vec]  # noqa: E731
     locals()
+
+
+@train_probe_ex.named_config
+def sum_distances():
+    attr_func = sum
+    attr_dim = 1
+    locals()

src/reward_preprocessing/scripts/gen_dots_and_dists.py

@@ -31,16 +31,22 @@ def generate_simple_trajectories(
     obs_list = []
     infos_list = []
     avg_distances = []
-    for i in range(num_transitions):
+    for i in range(num_transitions + 1):
         data, avg_distance, distances = generate_transition(
             number_pairs, circle_radius, size, colors, weights
         )
         obs_list.append(data)
         infos_list.append({"distances": distances})
         avg_distances.append(avg_distance)
 
-    # Duplicate last observation, since there is always a final next_obs.
-    obs_list.append(obs_list[-1].copy())
+    # Drop the first element of the infos list, since the first observation shouldn't
+    # come with an info dict (see flatten_trajectories_with_rew_double_info in
+    # common/utils.py)
+    infos_list = infos_list[1:]
+
+    # Drop the last element of avg_distances, since the last observation is the next_obs
+    # of the final transition, and rewards are associated with obs, not next_obs
+    avg_distances = avg_distances[:-1]
 
     condensed = {
         "obs": np.array(obs_list).astype(np.uint8),