def handshake(self): if self.client_state == States.CLOSED: # Step1: Generate a random sequence number. Send syn message to the server # to initiate a connection. seq_num = utils.rand_int() syn_header = utils.Header( seq_num, 0, syn=1, ) # for this case we send only header; # if you need to send data you will need to append it send_udp(syn_header.bits()) self.update_state(States.SYN_SENT) self.last_seq = seq_num self.handshake() elif self.client_state == States.SYN_SENT: # Step2: Rceive ack from server with syn = 1, ack = 1. header, body, addr = recv_msg() if header is None: print("Handshake failed.") self.update_state(States.CLOSED) elif header.ack == 1: seq_num = header.ack_num ack_num = header.seq_num + 1 syn_header = utils.Header(seq_num, ack_num, ack=1) send_udp(syn_header.bits()) self.last_seq = seq_num self.last_received_seq = header.seq_num self.last_received_ack = header.ack_num + 1 self.update_state(States.ESTABLISHED)
def handshake(self): while self.client_state != States.ESTABLISHED: if self.client_state == States.CLOSED: seq_num = utils.rand_int() self.next_seq_num = seq_num + 1 syn_header = utils.Header(seq_num, 0, syn=1, ack=0, fin=0) # for this case we send only header; # if you need to send data you will need to append it send_udp(syn_header.bits()) self.update_state(States.SYN_SENT) elif self.client_state == States.SYN_SENT: recv_data, addr = sock.recvfrom(1024) syn_ack_header = utils.bits_to_header(recv_data) self.last_received_seq_num = syn_ack_header.seq_num if syn_ack_header.syn == 1 and syn_ack_header.ack == 1: ack_header = utils.Header(self.next_seq_num, self.last_received_seq_num + 1, syn=0, ack=1, fin=0) self.next_seq_num += 1 send_udp(ack_header.bits()) self.update_state(States.ESTABLISHED) else: pass
def __getitem__(self, idx): args = self.args # args.time_step - number of time steps per episode # Offset is the last index in the episode that we can start at offset = args.time_step - args.n_his - args.n_roll + 1 src_rollout = idx // offset src_timestep = idx % offset ''' used for dynamics modeling ''' #TODO validate that this makes sense # load images for graph inference infer_st_idx = rand_int(0, args.time_step - args.n_identify + 1) # if using detected keypoints kps_pred = self.offline_data["observations"][(args.time_step)*src_rollout:args.time_step*(src_rollout+1)][::args.frame_offset] kps_preload_temp = np.concatenate([ kps_pred[infer_st_idx : infer_st_idx + args.n_identify], kps_pred[src_timestep : src_timestep + args.n_his + args.n_roll]], 0) kps_preload = np.zeros(shape = (kps_preload_temp.shape[0], args.n_kp, args.state_dim)) kps_preload[:, 0,0] = kps_preload_temp[:, 8] kps_preload[:, 0,1] = kps_preload_temp[:, 9] kps_preload[:, 1,0] = kps_preload_temp[:, 10] kps_preload[:, 1,1] = kps_preload_temp[:, 11] kps_preload[:, 2,0] = kps_preload_temp[:, 12] kps_preload[:, 2,1] = kps_preload_temp[:, 13] kps_preload[:, 3,0] = kps_preload_temp[:, 14] kps_preload[:, 3,1] = kps_preload_temp[:, 15] kps_preload[:, 4,0] = kps_preload_temp[:, 16] kps_preload[:, 4,1] = kps_preload_temp[:, 17] kps_preload[:, 5,0] = kps_preload_temp[:, 18] kps_preload[:, 5,1] = kps_preload_temp[:, 19] kps_preload = torch.FloatTensor(kps_preload) # get action actions_raw = self.offline_data["actions"][(args.time_step)*src_rollout:args.time_step*(src_rollout+1)][::args.frame_offset] actions_id = actions_raw[infer_st_idx:infer_st_idx + args.n_identify] actions_dy = actions_raw[src_timestep:src_timestep + args.n_his + args.n_roll] actions = np.concatenate([actions_id, actions_dy], 0) # if using preloaded keypoints return kps_preload, actions, self.node_params
def add_rels(self, param_load=None): param = np.zeros((self.n_ball * (self.n_ball - 1) // 2, 2)) self.param_dim = param.shape[0] if param_load is not None: print("Load param for init env") cnt = 0 rels_idx = [] for i in range(self.n_ball): for j in range(i): rel_type = rand_int( 0, self.n_rel_type) if param_load is None else param_load[cnt, 0] param[cnt, 0] = rel_type rels_idx.append([i, j]) pos_i = self.balls[i].position pos_j = self.balls[j].position if rel_type == 0: # no relation pass elif rel_type == 1: # spring rest_length = rand_float( 20, 120) if param_load is None else param_load[cnt, 1] param[cnt, 1] = rest_length c = pymunk.DampedSpring(self.balls[i], self.balls[j], (0, 0), (0, 0), rest_length=rest_length, stiffness=20, damping=0.) self.space.add(c) elif rel_type == 2: # string rest_length = calc_dis( pos_i, pos_j) if param_load is None else param_load[cnt, 1] param[cnt, 1] = rest_length c = pymunk.SlideJoint(self.balls[i], self.balls[j], (0, 0), (0, 0), rest_length - 5, rest_length + 5) self.space.add(c) else: raise AssertionError("Unknown relation type") cnt += 1 if param_load is not None: assert ((param == param_load).all()) self.rels_idx = rels_idx self.param = param
def handshake(self): if self.client_state == States.CLOSED: seq_num = utils.rand_int() syn_header = utils.Header(seq_num, 0, syn = 1, ack = 0) # for this case we send only header; # if you need to send data you will need to append it send_udp(syn_header.bits()) self.update_state(States.SYN_SENT) else: pass
def reg_room(cls, room: 'Room') -> 'Room': if room.id is not None: raise AssertionError latest_room: list = list(cls.rooms.keys())[-1:] if not latest_room: alloc_room_id = rand_int() else: alloc_room_id = cls.rooms[latest_room[0]].id + 1 room.id = alloc_room_id cls.rooms[str(room.id)] = room return room
def get_crop_params(phase, img, crop_size): w, h = img.size if w < h: tw = crop_size th = int(crop_size * h / w) else: th = crop_size tw = int(crop_size * w / h) if phase == 'train': if w == tw and h == th: return 0, 0, h, w assert False i = rand_int(0, h - th) j = rand_int(0, w - tw) else: i = int(round((h - th) / 2.)) j = int(round((w - tw) / 2.)) return i, j, th, tw
def __getitem__(self, idx): idx_rollout = idx // self.samples_per_rollout + self.n_rollout * self.idx idx_timestep = idx % self.samples_per_rollout # prepare input data seq_data = load_data(self.prepared_names, os.path.join(self.mother.data_dir, str(idx_rollout) + '.rollout.h5')) seq_data = [d[idx_timestep:idx_timestep + args.len_seq + 1] for d in seq_data] # prepare fit data fit_idx = rand_int(0, args.group_size - 1) # new traj idx in group fit_idx = fit_idx + idx_rollout // args.group_size * args.group_size # new traj idx in global fit_data = load_data(self.prepared_names, os.path.join(self.mother.data_dir, str(fit_idx) + '.rollout.h5')) return seq_data, fit_data
def handshake(self): seq_num = 0 if self.client_state == States.CLOSED: seq_num = utils.rand_int() syn_header = utils.Header(seq_num, 0, 1, 0, 0) # for this case we send only header; # if you need to send data you will need to append it send_udp(syn_header.bits()) self.update_state(States.SYN_SENT) elif self.client_state == States.SYN_SENT: header, body, addr = recv_msg() if header.syn == 1 and header.ack == 1 and header.ack_num == seq_num + 1: ack_num = header.seq_num + 1 seq_num = header.ack_num header = utils.Header(seq_num, ack_num, 0, 1, 0) send_udp(header.bits()) self.update_state(States.ESTABLISHED)
def send_reliable_message(self, message): count= 0 MSS = 2 # split the message into chunks msg = message.split() #seq number seq_num = utils.rand_int() # Create the header we attach to our message msg_header = utils.Header(seq_num, 0, syn = 0, ack = 1, fin = 0) # Send initial chunk chunk = msg[count:count + 3] chunk = ','.join(chunk[0:2]) send_udp(msg_header.bits() + chunk.replace(",", " ").encode()) count += 2 while count <= 7: #handle seq / ack recv_data, addr = sock.recvfrom(1024) #receive the data sent from the server msg_header = utils.bits_to_header(recv_data) #convert received data to header format firstseq = msg_header.seq_num # first msg seq number if msg_header.seq_num == firstseq or msg_header.seq_num == subseq: #check current seq numb from previous, make sure increments by 1 chunk = msg[MSS:MSS + 3] # chunk = ','.join(chunk[0:2]) send_udp(msg_header.bits() + chunk.replace(",", " ").encode()) count += 2 MSS += 2 subseq = msg_header.seq_num + 1 else: #if seq are not incrementing correctly then send previous chunk again print(msg_header.seq_num) #just printing out the seq number MSS -= 2 chunk = msg[MSS:MSS - 3] # chunk = ','.join(chunk[0:2]) send_udp(msg_header.bits() + chunk.replace(",", " ").encode()) #handle stop and wait pass
def terminate(self): while True: if self.client_state == States.ESTABLISHED: seq_num = utils.rand_int() header = utils.Header(seq_num, 0, 0, 0, 1) send_udp(header.bits()) self.update_state(States.FIN_WAIT_1) elif self.client_state == States.FIN_WAIT_1: header, body, addr = recv_msg() if header.ack == 1 and header.ack_num == seq_num + 1: self.update_state(States.FIN_WAIT_2) elif self.client_state == States.FIN_WAIT_2: header, body, addr = recv_msg() if header.fin == 1 and header.ack == 1 and header.ack_num == seq_num + 1: seq_num = header.ack_num ack_num = header.seq_num + 1 header = utils.Header(seq_num, ack_num, 0, 1, 0) send_udp(header.bits()) # wait for 2MSL time.sleep(1) self.update_state(States.CLOSED)
def handshake(self): if self.client_state == States.CLOSED: seq_num = utils.rand_int() syn_header = utils.Header(seq_num, 0, syn=1, ack=0, fin=0) # for this case we send only header; # if you need to send data you will need to append it send_udp(syn_header.bits()) # update client seq number self.my_next_seq = seq_num + 1 self.update_state(States.SYN_SENT) else: raise RuntimeError("invalid states for start a handshake.") # we wait for server to send back a SYN-ACK message if self.client_state == States.SYN_SENT: recv_data, addr = sock.recvfrom(1024) header = utils.bits_to_header(recv_data) # server syn header should have both syn and ack fields # and the ack_num should be client next sequence number if header.ack == 1 and header.syn == 1 and header.ack_num == self.my_next_seq: self.server_next_seq = header.seq_num + 1 self.last_received_ack = header.ack_num + 1 else: raise RuntimeError( "invalid server SYN-reply, handshake failed.") # we send back ACK message ack_header = utils.Header(self.my_next_seq, self.server_next_seq, syn=0, ack=1, fin=0) send_udp(ack_header.bits()) # update my seq self.my_next_seq += 1 # update state -> connection established on the client's perspective self.update_state(States.ESTABLISHED) else: raise RuntimeError("invalid states for waiting server SYN-reply.")
# action based on current state and updates its state # accordingly # You will need to add more states, please update the possible # states in utils.py file while True: if server_state == States.CLOSED: # we already started listening, just update the state update_server_state(States.LISTEN) elif server_state == States.LISTEN: # we are waiting for a message header, body, addr = recv_msg() # if received message is a syn message, it's a connection # initiation if header.syn == 1: # update seq numbers my_next_seq = utils.rand_int( ) # we randomly pick a sequence number client_next_seq = header.seq_num + 1 # update client addr client_addr = addr # make a syn-ack header and send to to client syn_ack_header = utils.Header(my_next_seq, client_next_seq, syn=1, ack=1, fin=0) sock.sendto(syn_ack_header.bits(), client_addr) # update my seq number my_next_seq += 1 # update server state update_server_state(States.SYN_RECEIVED)
def gen_PyFleX(info): env, root_num = info['env'], info['root_num'] thread_idx, data_dir, data_names = info['thread_idx'], info['data_dir'], info['data_names'] n_rollout, n_instance = info['n_rollout'], info['n_instance'] time_step, time_step_clip = info['time_step'], info['time_step_clip'] shape_state_dim, dt = info['shape_state_dim'], info['dt'] env_idx = info['env_idx'] np.random.seed(round(time.time() * 1000 + thread_idx) % 2**32) ### NOTE: we might want to fix the seed for reproduction # positions, velocities if env_idx == 5: # RiceGrip stats = [init_stat(6), init_stat(6)] else: stats = [init_stat(3), init_stat(3)] import pyflex pyflex.init() for i in range(n_rollout): if i % 10 == 0: print("%d / %d" % (i, n_rollout)) rollout_idx = thread_idx * n_rollout + i rollout_dir = os.path.join(data_dir, str(rollout_idx)) os.system('mkdir -p ' + rollout_dir) if env == 'FluidFall': scene_params = np.zeros(1) pyflex.set_scene(env_idx, scene_params, thread_idx) n_particles = pyflex.get_n_particles() positions = np.zeros((time_step, n_particles, 3), dtype=np.float32) velocities = np.zeros((time_step, n_particles, 3), dtype=np.float32) for j in range(time_step_clip): p_clip = pyflex.get_positions().reshape(-1, 4)[:, :3] pyflex.step() for j in range(time_step): positions[j] = pyflex.get_positions().reshape(-1, 4)[:, :3] if j == 0: velocities[j] = (positions[j] - p_clip) / dt else: velocities[j] = (positions[j] - positions[j - 1]) / dt pyflex.step() data = [positions[j], velocities[j]] store_data(data_names, data, os.path.join(rollout_dir, str(j) + '.h5')) elif env == 'BoxBath': # BoxBath scene_params = np.zeros(1) pyflex.set_scene(env_idx, scene_params, thread_idx) n_particles = pyflex.get_n_particles() positions = np.zeros((time_step, n_particles, 3), dtype=np.float32) velocities = np.zeros((time_step, n_particles, 3), dtype=np.float32) for j in range(time_step_clip): pyflex.step() p = pyflex.get_positions().reshape(-1, 4)[:64, :3] clusters = [] st_time = time.time() kmeans = MiniBatchKMeans(n_clusters=root_num[0][0], random_state=0).fit(p) # print('Time on kmeans', time.time() - st_time) clusters.append([[kmeans.labels_]]) # centers = kmeans.cluster_centers_ ref_rigid = p for j in range(time_step): positions[j] = pyflex.get_positions().reshape(-1, 4)[:, :3] # apply rigid projection to ground truth XX = ref_rigid YY = positions[j, :64] # print("MSE init", np.mean(np.square(XX - YY))) X = XX.copy().T Y = YY.copy().T mean_X = np.mean(X, 1, keepdims=True) mean_Y = np.mean(Y, 1, keepdims=True) X = X - mean_X Y = Y - mean_Y C = np.dot(X, Y.T) U, S, Vt = np.linalg.svd(C) D = np.eye(3) D[2, 2] = np.linalg.det(np.dot(Vt.T, U.T)) R = np.dot(Vt.T, np.dot(D, U.T)) t = mean_Y - np.dot(R, mean_X) YY_fitted = (np.dot(R, XX.T) + t).T # print("MSE fit", np.mean(np.square(YY_fitted - YY))) positions[j, :64] = YY_fitted if j > 0: velocities[j] = (positions[j] - positions[j - 1]) / dt pyflex.step() data = [positions[j], velocities[j], clusters] store_data(data_names, data, os.path.join(rollout_dir, str(j) + '.h5')) elif env == 'FluidShake': # if env is FluidShake height = 1.0 border = 0.025 dim_x = rand_int(10, 12) dim_y = rand_int(15, 20) dim_z = 3 x_center = rand_float(-0.2, 0.2) x = x_center - (dim_x-1)/2.*0.055 y = 0.055/2. + border + 0.01 z = 0. - (dim_z-1)/2.*0.055 box_dis_x = dim_x * 0.055 + rand_float(0., 0.3) box_dis_z = 0.2 scene_params = np.array([x, y, z, dim_x, dim_y, dim_z, box_dis_x, box_dis_z]) pyflex.set_scene(env_idx, scene_params, 0) boxes = calc_box_init_FluidShake(box_dis_x, box_dis_z, height, border) for i in range(len(boxes)): halfEdge = boxes[i][0] center = boxes[i][1] quat = boxes[i][2] pyflex.add_box(halfEdge, center, quat) n_particles = pyflex.get_n_particles() n_shapes = pyflex.get_n_shapes() # print("n_particles", n_particles) # print("n_shapes", n_shapes) positions = np.zeros((time_step, n_particles + n_shapes, 3), dtype=np.float32) velocities = np.zeros((time_step, n_particles + n_shapes, 3), dtype=np.float32) shape_quats = np.zeros((time_step, n_shapes, 4), dtype=np.float32) x_box = x_center v_box = 0. for j in range(time_step_clip): x_box_last = x_box x_box += v_box * dt shape_states_ = calc_shape_states_FluidShake( x_box, x_box_last, scene_params[-2:], height, border) pyflex.set_shape_states(shape_states_) pyflex.step() for j in range(time_step): x_box_last = x_box x_box += v_box * dt v_box += rand_float(-0.15, 0.15) - x_box * 0.1 shape_states_ = calc_shape_states_FluidShake( x_box, x_box_last, scene_params[-2:], height, border) pyflex.set_shape_states(shape_states_) positions[j, :n_particles] = pyflex.get_positions().reshape(-1, 4)[:, :3] shape_states = pyflex.get_shape_states().reshape(-1, shape_state_dim) for k in range(n_shapes): positions[j, n_particles + k] = shape_states[k, :3] shape_quats[j, k] = shape_states[k, 6:10] if j > 0: velocities[j] = (positions[j] - positions[j - 1]) / dt pyflex.step() # NOTE: 1) particle + glass wall positions, 2) particle + glass wall velocitys, 3) glass wall rotations, 4) scenen parameters data = [positions[j], velocities[j], shape_quats[j], scene_params] store_data(data_names, data, os.path.join(rollout_dir, str(j) + '.h5')) elif env == 'RiceGrip': # if env is RiceGrip # repeat the grip for R times R = 3 gripper_config = sample_control_RiceGrip() if i % R == 0: ### set scene # x, y, z: [8.0, 10.0] # clusterStiffness: [0.3, 0.7] # clusterPlasticThreshold: [0.00001, 0.0005] # clusterPlasticCreep: [0.1, 0.3] x = rand_float(8.0, 10.0) y = rand_float(8.0, 10.0) z = rand_float(8.0, 10.0) clusterStiffness = rand_float(0.3, 0.7) clusterPlasticThreshold = rand_float(0.00001, 0.0005) clusterPlasticCreep = rand_float(0.1, 0.3) scene_params = np.array([x, y, z, clusterStiffness, clusterPlasticThreshold, clusterPlasticCreep]) pyflex.set_scene(env_idx, scene_params, thread_idx) scene_params[4] *= 1000. halfEdge = np.array([0.15, 0.8, 0.15]) center = np.array([0., 0., 0.]) quat = np.array([1., 0., 0., 0.]) pyflex.add_box(halfEdge, center, quat) pyflex.add_box(halfEdge, center, quat) n_particles = pyflex.get_n_particles() n_shapes = pyflex.get_n_shapes() positions = np.zeros((time_step, n_particles + n_shapes, 6), dtype=np.float32) velocities = np.zeros((time_step, n_particles + n_shapes, 6), dtype=np.float32) shape_quats = np.zeros((time_step, n_shapes, 4), dtype=np.float32) for j in range(time_step_clip): shape_states = calc_shape_states_RiceGrip(0, dt, shape_state_dim, gripper_config) pyflex.set_shape_states(shape_states) pyflex.step() p = pyflex.get_positions().reshape(-1, 4)[:, :3] clusters = [] st_time = time.time() kmeans = MiniBatchKMeans(n_clusters=root_num[0][0], random_state=0).fit(p) # print('Time on kmeans', time.time() - st_time) clusters.append([[kmeans.labels_]]) # centers = kmeans.cluster_centers_ for j in range(time_step): shape_states = calc_shape_states_RiceGrip(j * dt, dt, shape_state_dim, gripper_config) pyflex.set_shape_states(shape_states) positions[j, :n_particles, :3] = pyflex.get_rigidGlobalPositions().reshape(-1, 3) positions[j, :n_particles, 3:] = pyflex.get_positions().reshape(-1, 4)[:, :3] shape_states = pyflex.get_shape_states().reshape(-1, shape_state_dim) for k in range(n_shapes): positions[j, n_particles + k, :3] = shape_states[k, :3] positions[j, n_particles + k, 3:] = shape_states[k, :3] shape_quats[j, k] = shape_states[k, 6:10] if j > 0: velocities[j] = (positions[j] - positions[j - 1]) / dt pyflex.step() data = [positions[j], velocities[j], shape_quats[j], clusters, scene_params] store_data(data_names, data, os.path.join(rollout_dir, str(j) + '.h5')) else: raise AssertionError("Unsupported env") # change dtype for more accurate stat calculation # only normalize positions and velocities datas = [positions.astype(np.float64), velocities.astype(np.float64)] # NOTE: stats is of length 2, for positions and velocities for j in range(len(stats)): stat = init_stat(stats[j].shape[0]) stat[:, 0] = np.mean(datas[j], axis=(0, 1))[:] stat[:, 1] = np.std(datas[j], axis=(0, 1))[:] stat[:, 2] = datas[j].shape[0] * datas[j].shape[1] stats[j] = combine_stat(stats[j], stat) pyflex.clean() return stats
def main(): # init args = init_arg_parser() P.log_path = args.log_path P.op_num = args.op_num P.lock_type = C.LOCK_MAPPING[args.lock_type] P.consensus_type = C.CONSENSUS_MAPPING[args.consensus_type] P.follower_num = args.follower_num P.lock_key_num = args.lock_key_num F.check_path_validity() time_suffix = time.strftime('%y%m%d_%H%M%S') logger = init_logger(time_suffix=time_suffix) F.save_parameter(time_suffix=time_suffix) # init servers ports = F.rand_items_in_range(P.leader_num + P.follower_num, P.port_range) servers = [] for i in range(P.leader_num + P.follower_num): if i == 0: leader = Leader(ports[i]) servers.append(leader) threading.Thread(target=leader.init_server).start() else: follower = Follower(ports[i]) leader.add_follower(follower) follower.add_leader(leader) servers.append(follower) threading.Thread(target=follower.init_server).start() # init clients time.sleep(0.1) clients = [] for i in range(P.client_num): client = Client() client.add_corresponding_server(servers[i]) clients.append(client) # client requests lock_keys = F.rand_items_in_range(P.lock_key_num, P.lock_key_range) thread_list = [] for i in range(P.op_num): logger.info("Operation No.{}".format(i + 1)) client = F.rand_item(clients) operation_code = F.rand_int(P.client_operation_range) operation_name = C.OPERATION_MAPPING[operation_code] func = getattr(client, operation_name) t = threading.Thread(target=func, args=(str(F.rand_item(lock_keys)),)) thread_list.append(t) t.start() time.sleep(0.015) # join client requests for t in thread_list: t.join() # shut down servers time.sleep(0.1) for s in servers: logger.info('Shutting down server {}'.format(s.get_short_uuid())) client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) client_socket.connect((s.ip, s.port)) send_data = {'type': 'close'} encoded_send_data = json.dumps(send_data).encode('utf-8') client_socket.send(encoded_send_data) client_socket.close() # finish time.sleep(5) logger.info('All done!')
server_state = new_state def recv_msg(): data, addr = sock.recvfrom(1024) header = utils.bits_to_header(data) body = utils.get_body_from_data(data) return (header, body, addr) while True: if server_state == States.CLOSED: update_server_state(States.LISTEN) elif server_state == States.LISTEN: header, body, addr = recv_msg() last_received_seq_num = header.seq_num if header.syn == 1: seq_number = utils.rand_int() syn_ack_msg = utils.Header(seq_number, ack_num = last_received_seq_num + 1, syn = 1, ack = 1, fin = 0) seq_number += 1 sock.sendto(syn_ack_msg.bits(), addr) update_server_state(States.SYN_RECEIVED) elif server_state == States.SYN_RECEIVED: header, body, addr = recv_msg() last_received_seq_num = header.seq_num if header.ack == 1: update_server_state(States.ESTABLISHED) pass elif server_state == States.SYN_SENT: pass elif server_state == States.ESTABLISHED: header, body, addr = recv_msg() last_received_seq_num = header.seq_num
# (Batch_size x n_samples x n_keypoints x n_state_features) kps = kps.view(B, n_samples, n_kp, args.state_dim) kps_id, kps_dy = kps[:, :n_identify], kps[:, n_identify:] # only train dynamics module kps = kps.detach() actions_id, actions_dy = actions[:, : n_identify], actions[:, n_identify:] ''' step #1: identify the dynamics graph ''' # randomize the observation length observe_length = rand_int(args.min_res, n_identify + 1) if args.baseline == 1: graph = model_dy.init_graph(kps_id[:, :observe_length], use_gpu=True, hard=True) else: #TODO Make sure that the actions are correctly encoded in the edges graph = model_dy.graph_inference( kps_id[:, :observe_length], actions_id[:, :observe_length], env=args.env, node_params=node_params) # calculate edge calculation accuracy
def __getitem__(self, idx): args = self.args suffix = '.png' if args.env in ['Ball'] else '.jpg' if args.stage == 'kp': src_rollout = idx // args.time_step src_timestep = idx % args.time_step elif args.stage in 'dy': offset = args.time_step - args.n_his - args.n_roll + 1 src_rollout = idx // offset src_timestep = idx % offset ''' used for keypoint detection ''' if args.stage == 'kp': src_path = os.path.join(args.dataf, self.phase, str(src_rollout), 'fig_%d%s' % (src_timestep, suffix)) # use the same rollout if in Cloth # des_rollout = rand_int(0, self.n_rollout) if args.env in ['Ball'] else src_rollout des_rollout = rand_int(0, self.n_rollout) des_timestep = rand_int(0, args.time_step) des_path = os.path.join(args.dataf, self.phase, str(des_rollout), 'fig_%d%s' % (des_timestep, suffix)) src = self.loader(src_path) des = self.loader(des_path) src = resize_and_crop(self.phase, src, self.scale_size, self.crop_size) des = resize_and_crop(self.phase, des, self.scale_size, self.crop_size) src = self.trans_to_tensor(src) des = self.trans_to_tensor(des) return src, des ''' used for dynamics modeling ''' if args.stage in 'dy': imgs = [] kp_preload = None # load images for graph inference infer_st_idx = rand_int(0, args.time_step - args.n_identify + 1) # if using detected keypoints if args.preload_kp == 1: # if using preload keypoints path = os.path.join(args.dataf + '_nKp_%d' % args.n_kp, self.phase, str(src_rollout) + '.h5') kps_pred = load_data(['keypoints'], path)[0][::args.frame_offset] kps_preload = np.concatenate([ kps_pred[infer_st_idx : infer_st_idx + args.n_identify], kps_pred[src_timestep : src_timestep + args.n_his + args.n_roll]], 0) kps_preload = torch.FloatTensor(kps_preload) else: # if detect keypoints during runtime for i in range(infer_st_idx, infer_st_idx + args.n_identify): path = os.path.join(args.dataf, self.phase, str(src_rollout), 'fig_%d%s' % (i, suffix)) img = self.loader(path) img = resize_and_crop(self.phase, img, self.scale_size, self.crop_size) img = self.trans_to_tensor(img) imgs.append(img) # load images for dynamics prediction for i in range(args.n_his + args.n_roll): path = os.path.join(args.dataf, self.phase, str(src_rollout), 'fig_%d%s' % (src_timestep + i, suffix)) img = self.loader(path) img = resize_and_crop(self.phase, img, self.scale_size, self.crop_size) img = self.trans_to_tensor(img) imgs.append(img) imgs = torch.cat(imgs, 0) assert imgs.size(0) == (args.n_identify + args.n_his + args.n_roll) * 3 if args.env in ['Ball']: # get ground truth edge type data_path = os.path.join(args.dataf, self.phase, str(src_rollout) + '.h5') metadata = load_data(self.data_names, data_path) edge_type = metadata[3][0, :, 0].astype(np.int) edge_attr = metadata[3][0, :, 1:] edge_type_gt = np.zeros((args.n_kp, args.n_kp, args.edge_type_num)) edge_attr_gt = np.zeros((args.n_kp, args.n_kp, edge_attr.shape[1])) cnt = 0 for x in range(args.n_kp): for y in range(x): edge_type_gt[x, y, edge_type[cnt]] = 1. edge_type_gt[y, x, edge_type[cnt]] = 1. edge_attr_gt[x, y] = edge_attr[cnt] edge_attr_gt[y, x] = edge_attr[cnt] cnt += 1 edge_type_gt = torch.FloatTensor(edge_type_gt) edge_attr_gt = torch.FloatTensor(edge_attr_gt) graph_gt = edge_type_gt, edge_attr_gt # get ground truth keypoint position states = metadata[1] / 80. kps_gt_id = states[infer_st_idx:infer_st_idx + args.n_identify, :, :2] kps_gt_dy = states[src_timestep:src_timestep + args.n_his + args.n_roll, :, :2] kps_gt = np.concatenate([kps_gt_id, kps_gt_dy], 0) kps_gt[:, :, 1] *= -1 kps_gt = torch.FloatTensor(kps_gt) actions = metadata[2] / 600. actions_id = actions[infer_st_idx:infer_st_idx + args.n_identify] actions_dy = actions[src_timestep:src_timestep + args.n_his + args.n_roll] actions = np.concatenate([actions_id, actions_dy], 0) actions = torch.FloatTensor(actions) # actions: (n_identify + n_his + n_roll) x n_kp x action_dim # print('actions size', actions.size()) # if using detected keypoints if args.preload_kp == 1: # if using preloaded keypoints return kps_preload, kps_gt, graph_gt, actions else: # if detecting keypoints during runtime return imgs, kps_gt, graph_gt, actions elif args.env in ['Cloth']: # get action data_path = os.path.join(args.dataf, self.phase, str(src_rollout) + '.h5') metadata = load_data(self.data_names, data_path) states = metadata[0][::args.frame_offset] actions_raw = metadata[1][::args.frame_offset] scene_params = metadata[2] stiffness = scene_params[15] ctrl_idx = scene_params[7:15].astype(np.int) states_id = states[infer_st_idx:infer_st_idx + args.n_identify] states_dy = states[src_timestep:src_timestep + args.n_his + args.n_roll] actions_id_raw = actions_raw[infer_st_idx:infer_st_idx + args.n_identify] actions_dy_raw = actions_raw[src_timestep:src_timestep + args.n_his + args.n_roll] # generate actions_id / actions_dy actions_id = np.zeros((args.n_identify, 6)) actions_dy = np.zeros((args.n_his + args.n_roll, 6)) actions_id[:, :3] = states_id[ np.arange(actions_id.shape[0]), ctrl_idx[actions_id_raw[:, 0, 0].astype(np.int)], :3] / 0.5 # normalize actions_dy[:, :3] = states_dy[ np.arange(actions_dy.shape[0]), ctrl_idx[actions_dy_raw[:, 0, 0].astype(np.int)], :3] / 0.5 # normalize actions_id[:, 3:] = actions_id_raw[:, 0, 1:] / 0.03 # normalize actions_dy[:, 3:] = actions_dy_raw[:, 0, 1:] / 0.03 # normalize actions_id = torch.FloatTensor(actions_id)[:, None, :].repeat(1, args.n_kp, 1) actions_dy = torch.FloatTensor(actions_dy)[:, None, :].repeat(1, args.n_kp, 1) actions = torch.cat([actions_id, actions_dy], 0) # if using detected keypoints if args.preload_kp == 1: # if using preloaded keypoints return kps_preload, actions else: # if detecting keypoints during runtime return imgs, actions else: raise AssertionError("Unknown env %s" % args.env)
def gen_Cloth(info): env, env_idx = info['env'], info['env_idx'] thread_idx, data_dir, data_names = info['thread_idx'], info['data_dir'], info['data_names'] n_rollout, time_step = info['n_rollout'], info['time_step'] dt, args, phase = info['dt'], info['args'], info['phase'] vis_width, vis_height = info['vis_width'], info['vis_height'] state_dim = args.state_dim action_dim = args.action_dim dt = 1. / 60. np.random.seed(round(time.time() * 1000 + thread_idx) % 2 ** 32) stats = [init_stat(state_dim), init_stat(action_dim)] engine = ClothEngine(dt, state_dim, action_dim) import pyflex pyflex.init() # bar = ProgressBar() for i in range(n_rollout): rollout_idx = thread_idx * n_rollout + i rollout_dir = os.path.join(data_dir, str(rollout_idx)) os.system('mkdir -p ' + rollout_dir) engine.init(pyflex) scene_params = engine.scene_params action = np.zeros(4) states_all = np.zeros((time_step, engine.n_particles, state_dim)) actions_all = np.zeros((time_step, 1, action_dim)) # drop the cloth down engine.set_action(action) engine.step() for j in range(time_step): positions = pyflex.get_positions().reshape(-1, 4)[:, :3] # sample the action if j % 5 == 0: ctrl_pts = rand_int(0, 8) act_lim = 0.05 dx = rand_float(-act_lim, act_lim) dz = rand_float(-act_lim, act_lim) dy = 0.05 action = np.array([ctrl_pts, dx, dy, dz]) else: action[2] = 0. # store the rollout information state = engine.get_state() states_all[j] = state tga_path = os.path.join(rollout_dir, '%d.tga' % j) pyflex.render(capture=True, path=tga_path) tga = Image.open(tga_path) tga = np.array(tga)[:, 60:780, :3][:, :, ::-1] tga = cv2.resize(tga, (vis_width, vis_height), interpolation=cv2.INTER_AREA) os.system('rm ' + tga_path) jpg_path = os.path.join(rollout_dir, 'fig_%d.jpg' % j) cv2.imwrite(jpg_path, tga) actions_all[j, 0] = action.copy() engine.set_action(action) engine.step() datas = [states_all, actions_all, scene_params] store_data(data_names, datas, rollout_dir + '.h5') datas = [datas[j].astype(np.float64) for j in range(len(datas))] for j in range(len(stats)): stat = init_stat(stats[j].shape[0]) stat[:, 0] = np.mean(datas[j], axis=(0, 1))[:] stat[:, 1] = np.std(datas[j], axis=(0, 1))[:] stat[:, 2] = datas[j].shape[0] stats[j] = combine_stat(stats[j], stat) pyflex.clean() return stats
def init(self, pyflex, scene_params=None): self.pyflex = pyflex if scene_params is None: # offset radius = 0.05 offset_x = -1. offset_y = 0.06 offset_z = -1. # fabrics fabric_type = rand_int(0, 3) # 0: Cloth, 1: shirt, 2: pants if fabric_type == 0: # parameters of the shape dimx = rand_int(25, 35) # dimx, width dimy = rand_int(25, 35) # dimy, height dimz = 0 # the actuated points ctrl_idx = np.array([ 0, dimx // 2, dimx - 1, dimy // 2 * dimx, dimy // 2 * dimx + dimx - 1, (dimy - 1) * dimx, (dimy - 1) * dimx + dimx // 2, (dimy - 1) * dimx + dimx - 1 ]) offset_x = -dimx * radius / 2. offset_y = 0.06 offset_z = -dimy * radius / 2. elif fabric_type == 1: # parameters of the shape dimx = rand_int(16, 25) # width of the body dimy = rand_int(30, 35) # height of the body dimz = 7 # size of the sleeves # the actuated points ctrl_idx = np.array([ dimx * dimy, dimx * dimy + dimz * (dimz + dimz // 2) + (1 + dimz) * (dimz + 1) // 4, dimx * dimy + (1 + dimz) * dimz // 2 + dimz * (dimz - 1), dimx * dimy + dimz * (dimz + dimz // 2) + (1 + dimz) * (dimz + 1) // 4 + \ (1 + dimz) * dimz // 2 + dimz * dimz - 1, dimy // 2 * dimx, dimy // 2 * dimx + dimx - 1, (dimy - 1) * dimx, dimy * dimx - 1]) offset_x = -(dimx + dimz * 4) * radius / 2. offset_y = 0.06 offset_z = -dimy * radius / 2. elif fabric_type == 2: # parameters of the shape dimx = rand_int(9, 13) * 2 # width of the pants dimy = rand_int(6, 11) # height of the top part dimz = rand_int(24, 31) # height of the leg # the actuated points ctrl_idx = np.array([ 0, dimx - 1, (dimy - 1) * dimx, (dimy - 1) * dimx + dimx - 1, dimx * dimy + dimz // 2 * (dimx - 4) // 2, dimx * dimy + (dimz - 1) * (dimx - 4) // 2, dimx * dimy + dimz * (dimx - 4) // 2 + 3 + \ dimz // 2 * (dimx - 4) // 2 + (dimx - 4) // 2 - 1, dimx * dimy + dimz * (dimx - 4) // 2 + 3 + \ dimz * (dimx - 4) // 2 - 1]) offset_x = -dimx * radius / 2. offset_y = 0.06 offset_z = -(dimy + dimz) * radius / 2. # physical param stiffness = rand_float(0.4, 1.0) stretchStiffness = stiffness bendStiffness = stiffness shearStiffness = stiffness dynamicFriction = 0.6 staticFriction = 1.0 particleFriction = 0.6 invMass = 1.0 # other parameters windStrength = 0.0 draw_mesh = 1. # set up environment self.scene_params = np.array([ offset_x, offset_y, offset_z, fabric_type, dimx, dimy, dimz, ctrl_idx[0], ctrl_idx[1], ctrl_idx[2], ctrl_idx[3], ctrl_idx[4], ctrl_idx[5], ctrl_idx[6], ctrl_idx[7], stretchStiffness, bendStiffness, shearStiffness, dynamicFriction, staticFriction, particleFriction, invMass, windStrength, draw_mesh ]) else: self.scene_params = scene_params scene_idx = 15 self.pyflex.set_scene(scene_idx, self.scene_params, 0) # set up camera pose camPos = np.array([0., 3.5, 0.]) camAngle = np.array([0., -90. / 180. * np.pi, 0.]) pyflex.set_camPos(camPos) pyflex.set_camAngle(camAngle) self.n_particles = self.pyflex.get_n_particles() # let the cloth drop action_zero = np.zeros(4) for i in range(5): pyflex.step(action_zero)
# action based on current state and updates its state # accordingly # You will need to add more states, please update the possible # states in utils.py file while True: if server_state == utils.States.CLOSED: # we already started listening, just update the state update_server_state(utils.States.LISTEN) elif server_state == utils.States.LISTEN: # we are waiting for a message header, body, addr = recv_msg() last_received_seq_num = header.seq_num # if received message is a syn message, it's a connection # initiation if header.syn == 1: seq_number = utils.rand_int() # we randomly pick a sequence number syn_ack_msg = utils.Header(seq_number, ack_num=last_received_seq_num + 1, syn=1, ack=1, fin=0) seq_number += 1 sock.sendto(syn_ack_msg.bits(), addr) update_server_state(utils.States.SYN_RECEIVED) ack_number = header.seq_num + 1 chunk = utils # to be implemented ### sending message from the server: # use the following method to send messages back to client
def gen_args(): args = parser.parse_args() if args.env == 'Ball': args.data_names = ['attrs', 'states', 'actions', 'rels'] args.frame_offset = 1 args.train_valid_ratio = 0.99 # radius args.attr_dim = 1 # x, y, xdot, ydot args.state_dim = 4 # ddx, ddy args.action_dim = 2 # none, spring, rod args.relation_dim = 3 # size of the latent causal graph args.node_attr_dim = 0 args.edge_attr_dim = 1 args.edge_type_num = 3 # generate random relations # added this nb_edges = args.n_ball * (args.n_ball - 1) // 2 load_rels = np.zeros((nb_edges, 2)) # the balls that are connected have the same type of relation (spring or string) and the same attribute for i in range(nb_edges): if rand_int(0, 2) == 1: load_rels[i, 0] = args.rel_type load_rels[i, 1] = args.rel_attr args.load_rels = load_rels args.height_raw = 110 args.width_raw = 110 args.height = 64 args.width = 64 args.scale_size = 64 args.crop_size = 64 args.lim = [-1., 1., -1., 1.] elif args.env == 'Cloth': args.data_names = ['states', 'actions', 'scene_params'] args.n_rollout = 2000 if args.stage == 'dy': args.frame_offset = 5 else: args.frame_offset = 1 args.time_step = 300 // args.frame_offset args.train_valid_ratio = 0.9 # x, y, z, xdot, ydot, zdot args.state_dim = 6 # x, y, z, dx, dy, dz args.action_dim = 6 # size of the latent causal graph args.node_attr_dim = 0 args.edge_attr_dim = 1 args.edge_type_num = 2 args.height_raw = 400 args.width_raw = 400 args.height = 64 args.width = 64 args.scale_size = 64 args.crop_size = 64 args.lim = [-1., 1., -1., 1.] else: raise AssertionError("Unsupported env %s" % args.env) # path to data args.dataf = 'data/' + args.dataf + '_' + args.env return args