from prioritized_memory import Memory
import numpy as np
import cv2
from utils import Transition
R = 5.0
m1 = Memory(1000)
m2 = Memory(1000)
m3 = Memory(1000)
M1 = Memory(1500)
M2 = Memory(1500)
M3 = Memory(1500)
m1.load_memory("../training/logger_013/suction_1_memory.pkl")
m2.load_memory("../training/logger_013/suction_2_memory.pkl")
m3.load_memory("../training/logger_013/gripper_memory.pkl")
empty_color = []
empty_depth = []
for i in range(m1.length):
M1.add(m1.tree.data[i])
M2.add(m2.tree.data[i])
M3.add(m3.tree.data[i])
for i in range(m1.length):
# Invalid point is common
if m1.tree.data[i].reward == -3 * R:
transition = m1.tree.data[i]
pixel_index = transition.pixel_idx
workspace_limits = np.asarray([[x_lower, x_upper], [y_lower, y_upper], [z_lower, z_upper]]) # X Y Z
heightmap_resolution = (workspace_limits[0][1]-workspace_limits[0][0])/resolution
# Constant
reward = 5.0
discount_factor = 0.5
iteration = 0
learned_times = 0
sufficient_exp = 0
buffer_size = 500
gripper_memory = Memory(buffer_size)
suction_1_memory = Memory(buffer_size)
suction_2_memory = Memory(buffer_size)
gripper_memory.load_memory("../training/mannual_label/gripper_memory.pkl")
suction_1_memory.load_memory("../training/mannual_label/suction_1_memory.pkl")
suction_2_memory.load_memory("../training/manual_label/suction_2_memory.pkl")
# Service clients
vacuum_pump_control = rospy.ServiceProxy("/vacuum_pump_control_node/vacuum_control", SetBool)
check_suck_success = rospy.ServiceProxy("/vacuum_pump_control_node/check_suck_success", SetBool)
agent_take_action_client = rospy.ServiceProxy("/agent_server_node/agent_take_action", agent_abb_action)
calibrate_gripper_client = rospy.ServiceProxy("/change_tool_service/calibrate_gripper", SetBool)
get_pc_client = rospy.ServiceProxy("/combine_pc_node/get_pc", get_pc)
empty_checker = rospy.ServiceProxy("/combine_pc_node/empty_state", pc_is_empty)
check_grasp_success = rospy.ServiceProxy("/combine_pc_node/grasp_state", check_grasp_success)
go_home = rospy.ServiceProxy("/agent_server_node/go_home", Empty)
go_place = rospy.ServiceProxy("/agent_server_node/go_place", Empty)
fixed_home = rospy.ServiceProxy("/agent_server_node/go_home_fix_orientation", Empty)
check_if_collide_client = rospy.ServiceProxy("/agent_server_node/check_if_collide", agent_abb_action)
train_iter = 30
save_freq = 5
save_root = "models_with_novel_422_batch_{}".format(mini_batch_size)
primitive_lr = 2.5e-4
dexnet_lr = 5e-5
suction_1_sampled = np.zeros(memory_capacity[0])
suction_2_sampled = np.zeros(memory_capacity[1])
gripper_sampled = np.zeros(memory_capacity[2])
#model_str = "../training/logger_010/models/10_121.pth"
model_str = "evaluate_model/600.pth"
suction_1_memory = Memory(memory_capacity[0])
suction_2_memory = Memory(memory_capacity[1])
gripper_memory = Memory(memory_capacity[2])
suction_1_memory.load_memory("suction_1_mixed_400.pkl")
suction_2_memory.load_memory("suction_2_mixed_200.pkl")
gripper_memory.load_memory("gripper_mixed_200.pkl")
#suction_1_memory.tree.reset_priority()
#suction_2_memory.tree.reset_priority()
#gripper_memory.tree.reset_priority()
compare_color = "../training/logger_009/images/color_000005.jpg"
compare_depth = "../training/logger_009/depth_data/depth_data_000005.npy"
def create_path():
cwd = os.getcwd() + "/" + save_root
if not os.path.exists(cwd):
os.makedirs(cwd)
primitives = [
if model_str == "" and testing: # TEST SHOULD PROVIDE MODEL
print "\033[0;31mNo model provided, exit!\033[0m"
os._exit(0)
# Get logger path
r = rospkg.RosPack()
package_path = r.get_path("grasp_suck")
csv_path, image_path, depth_path, mixed_paths, feat_paths, pc_path, model_path, vis_path, diff_path, check_grasp_path = utils.getLoggerPath(
testing, package_path, run)
'''if run!=0:
gripper_memory_buffer.load_memory(csv_path.replace("logger_{:03}".format(run), "logger_{:03}".format(run-1))+"gripper_memory.pkl")
suction_1_memory_buffer.load_memory(csv_path.replace("logger_{:03}".format(run), "logger_{:03}".format(run-1))+"suction_1_memory.pkl")
suction_2_memory_buffer.load_memory(csv_path.replace("logger_{:03}".format(run), "logger_{:03}".format(run-1))+"suction_2_memory.pkl")
'''
if suction_1_memory != "":
suction_1_memory_buffer.load_memory(suction_1_memory)
if suction_2_memory != "":
suction_2_memory_buffer.load_memory(suction_2_memory)
if gripper_memory != "":
gripper_memory_buffer.load_memory(gripper_memory)
# trainer
if testing:
learning_rate = 5e-5
densenet_lr = 1e-5
trainer = Trainer(reward, discount_factor, use_cpu, learning_rate, densenet_lr)
# Load model if provided
if model_str != "":
print "[%f]: Loading provided model..." % (time.time())
trainer.behavior_net.load_state_dict(torch.load(model_str))
