-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathhelper.py
107 lines (88 loc) · 4.49 KB
/
helper.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
import numpy as np
import matplotlib.pyplot as plt
from constants import PICK_TARGETS, PLACE_TARGETS, PIXEL_SIZE, BOUNDS
from heapq import nlargest
#@title Helper Functions
def build_scene_description(found_objects, block_name="box", bowl_name="circle"):
scene_description = f"objects = {found_objects}"
scene_description = scene_description.replace(block_name, "block")
scene_description = scene_description.replace(bowl_name, "bowl")
scene_description = scene_description.replace("'", "")
return scene_description
def step_to_nlp(step):
step = step.replace("robot.pick_and_place(", "")
step = step.replace(")", "")
pick, place = step.split(", ")
return "Pick the " + pick + " and place it on the " + place + "."
def normalize_scores(scores):
max_score = max(scores.values())
normed_scores = {key: np.clip(scores[key] / max_score, 0, 1) for key in scores}
return normed_scores
def plot_saycan(llm_scores, vfs, combined_scores, task, correct=True, show_top=None):
if show_top:
top_options = nlargest(show_top, combined_scores, key = combined_scores.get)
# add a few top llm options in if not already shown
top_llm_options = nlargest(show_top // 2, llm_scores, key = llm_scores.get)
for llm_option in top_llm_options:
if not llm_option in top_options:
top_options.append(llm_option)
llm_scores = {option: llm_scores[option] for option in top_options}
vfs = {option: vfs[option] for option in top_options}
combined_scores = {option: combined_scores[option] for option in top_options}
sorted_keys = dict(sorted(combined_scores.items()))
keys = [key for key in sorted_keys]
positions = np.arange(len(combined_scores.items()))
width = 0.3
fig = plt.figure(figsize=(12, 6))
ax1 = fig.add_subplot(1,1,1)
plot_llm_scores = normalize_scores({key: np.exp(llm_scores[key]) for key in sorted_keys})
plot_llm_scores = np.asarray([plot_llm_scores[key] for key in sorted_keys])
plot_affordance_scores = np.asarray([vfs[key] for key in sorted_keys])
plot_combined_scores = np.asarray([combined_scores[key] for key in sorted_keys])
ax1.bar(positions, plot_combined_scores, 3 * width, alpha=0.6, color="#93CE8E", label="combined")
score_colors = ["#ea9999ff" for score in plot_affordance_scores]
ax1.bar(positions + width / 2, 0 * plot_combined_scores, width, color="#ea9999ff", label="vfs")
ax1.bar(positions + width / 2, 0 * plot_combined_scores, width, color="#a4c2f4ff", label="language")
ax1.bar(positions - width / 2, np.abs(plot_affordance_scores), width, color=score_colors)
plt.xticks(rotation="vertical")
ax1.set_ylim(0.0, 1.0)
ax1.grid(True, which="both")
ax1.axis("on")
ax1_llm = ax1.twinx()
ax1_llm.bar(positions + width / 2, plot_llm_scores, width, color="#a4c2f4ff", label="language")
ax1_llm.set_ylim(0.01, 1.0)
plt.yscale("log")
font = {"size":"16", "color":"k" if correct else "r"}
# font = {"fontname":"Arial", "size":"16", "color":"k" if correct else "r"}
plt.title(task, **font)
key_strings = [key.replace("robot.pick_and_place", "").replace(", ", " to ").replace("(", "").replace(")","") for key in keys]
plt.xticks(positions, key_strings, **font)
ax1.legend()
plt.show()
#@title Affordance Scoring
#@markdown Given this environment does not have RL-trained policies or an asscociated value function, we use affordances through an object detector.
def affordance_scoring(options, found_objects, verbose=False, block_name="box", bowl_name="circle", termination_string="done()"):
affordance_scores = {}
found_objects = [
found_object.replace(block_name, "block").replace(bowl_name, "bowl")
for found_object in found_objects + list(PLACE_TARGETS.keys())[-5:]]
verbose and print("found_objects", found_objects)
for option in options:
if option == termination_string:
affordance_scores[option] = 0.2
continue
pick, place = option.replace("robot.pick_and_place(", "").replace(")", "").split(", ")
affordance = 0
found_objects_copy = found_objects.copy()
if pick in found_objects_copy:
found_objects_copy.remove(pick)
if place in found_objects_copy:
affordance = 1
affordance_scores[option] = affordance
verbose and print(affordance, '\t', option)
return affordance_scores
def xyz_to_pix(position, bounds=BOUNDS, pixel_size=PIXEL_SIZE):
"""Convert from 3D position to pixel location on heightmap."""
u = int(np.round((position[1] - bounds[1, 0]) / pixel_size))
v = int(np.round((position[0] - bounds[0, 0]) / pixel_size))
return (u, v)