Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,688 Bytes
4f09ecf |
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 |
#######################################################################
# Name: robot.py
#
# - Stores S(t), A(t), R(t), S(t+1)
#######################################################################
from copy import deepcopy
import torch
class Robot:
def __init__(self, robot_id, position, plot=False):
self.robot_id = robot_id
self.plot = plot
self.travel_dist = 0
self.robot_position = position
self.observations = None
self.trajectory_coords = []
self.targets_found_on_path = []
self.episode_buffer = []
for i in range(15):
self.episode_buffer.append([])
if self.plot:
# initialize the route
self.xPoints = [self.robot_position[0]]
self.yPoints = [self.robot_position[1]]
def save_observations(self, observations):
node_inputs, edge_inputs, current_index, node_padding_mask, edge_padding_mask, edge_mask = observations
self.episode_buffer[0] += deepcopy(node_inputs).to('cpu')
self.episode_buffer[1] += deepcopy(edge_inputs).to('cpu')
self.episode_buffer[2] += deepcopy(current_index).to('cpu')
self.episode_buffer[3] += deepcopy(node_padding_mask).to('cpu')
self.episode_buffer[4] += deepcopy(edge_padding_mask).to('cpu')
self.episode_buffer[5] += deepcopy(edge_mask).to('cpu')
def save_action(self, action_index):
self.episode_buffer[6] += action_index.unsqueeze(0).unsqueeze(0)
def save_reward_done(self, reward, done):
self.episode_buffer[7] += deepcopy(torch.FloatTensor([[[reward]]])).to('cpu')
self.episode_buffer[8] += deepcopy(torch.tensor([[[(int(done))]]])).to('cpu')
if self.plot:
self.xPoints.append(self.robot_position[0])
self.yPoints.append(self.robot_position[1])
def save_next_observations(self, observations):
node_inputs, edge_inputs, current_index, node_padding_mask, edge_padding_mask, edge_mask = observations
self.episode_buffer[9] += deepcopy(node_inputs).to('cpu')
self.episode_buffer[10] += deepcopy(edge_inputs).to('cpu')
self.episode_buffer[11] += deepcopy(current_index).to('cpu')
self.episode_buffer[12] += deepcopy(node_padding_mask).to('cpu')
self.episode_buffer[13] += deepcopy(edge_padding_mask).to('cpu')
self.episode_buffer[14] += deepcopy(edge_mask).to('cpu')
# NEW: ADDED TO SAVE TRAJECTORY COORDS DURING TTA
def save_trajectory_coords(self, robot_position_coords, num_target_found):
self.trajectory_coords.append(robot_position_coords)
self.targets_found_on_path.append(num_target_found)
|