def __init__(self, config, robot, randomize=True):
self._physics_client = robot.physics_client
self._robot = robot
full_obs = config.get('full_observation', False)
intrinsics_path = config['camera_info']
extrinsics_path = config['transform']
extrinsics_dict = io_utils.load_yaml(extrinsics_path)
self._camera_info = io_utils.load_yaml(intrinsics_path)
self._transform = transform_utils.from_dict(extrinsics_dict)
self._randomize = config.get('randomize', None) if randomize else None
self._construct_camera(self._camera_info, self._transform)
self.state_space = gym.spaces.Box(low=0,
high=1,
shape=(self.camera.info.height,
self.camera.info.width, 1))
if full_obs:
#RGB + Depth
self.state_space = gym.spaces.Box(low=0,
high=255,
shape=(self.camera.info.height,
self.camera.info.width,
5))
def __init__(self, config, sensor, robot):
self.scope = 'encoded_img_sensor'
self._sensor = sensor
self._robot = robot
self.scene_type = config['scene'].get('scene_type', "OnTable")
config = config['sensor']
self._visualize = config.get('visualize', False)
# Load the encoder
model_dir = config['encoder_dir']
encoder_config = io_utils.load_yaml(
os.path.join(model_dir, 'config.yaml'))
# Build the graph and restore trained weights
with tf.name_scope(self.scope):
self._encoder = encoders.SimpleAutoEncoder(encoder_config)
self._encoder.load_weights(model_dir)
# Define the state space
dim = int(np.prod(self._encoder.encoding_shape))
self.state_space = gym.spaces.Box(-1., 1., (dim, ), np.float32)
# If requested, setup an OpenCV window for visualizations
if self._visualize:
cv2.namedWindow('imgs', flags=cv2.WINDOW_NORMAL)
def __init__(self, config, evaluate=False, test=False, validate=False):
if not isinstance(config, dict):
config = io_utils.load_yaml(config)
super().__init__(config, evaluate=evaluate, test=test, validate=validate)
self._step_time = collections.deque(maxlen=10000)
self.time_horizon = config['time_horizon']
self._workspace = {'lower': np.array([-1., -1., -1]),
'upper': np.array([1., 1., 1.])}
self.model_path = config['robot']['model_path']
self._simplified = config['simplified']
self.depth_obs = config.get('depth_observation', False)
self.full_obs = config.get('full_observation', False)
self._initial_height = 0.3
self._init_ori = transformations.quaternion_from_euler(np.pi, 0., 0.)
self.main_joints = [0, 1, 2, 3] #FIXME make it better
self._left_finger_id = 7
self._right_finger_id = 9
self._fingers = [self._left_finger_id, self._right_finger_id]
self._model = None
self._joints = None
self._left_finger, self._right_finger = None, None
self._actuator = actuator.Actuator(self, config, self._simplified)
self._camera = sensor.RGBDSensor(config['sensor'], self)
# Assign the reward fn
if self._simplified:
self._reward_fn = SimplifiedReward(config['reward'], self)
elif config['reward']['custom']:
self._reward_fn = ShapedCustomReward(config['reward'], self)
else:
self._reward_fn = Reward(config['reward'], self)
# Assign the sensors
if self.depth_obs or self.full_obs:
self._sensors = [self._camera]
else:
self._encoder = sensor.EncodedDepthImgSensor(
config, self._camera, self)
self._sensors = [self._encoder]
if not self._simplified:
self._sensors.append(self._actuator)
self.curriculum = WorkspaceCurriculum(config['curriculum'], self, evaluate)
self.history = self.curriculum._history
self._callbacks = {RobotEnv.Events.START_OF_EPISODE: [],
RobotEnv.Events.END_OF_EPISODE: [],
RobotEnv.Events.CLOSE: [],
RobotEnv.Events.CHECKPOINT: []}
self.register_events(evaluate, config)
self.sr_mean = 0.
self.setup_spaces()
def visualize(args):
n_imgs = 2 # number of images to visualize
# Load the model
config = io_utils.load_yaml(os.path.join(args.model_dir, 'config.yaml'))
model = encoders.SimpleAutoEncoder(config)
model.load_weights(args.model_dir)
# Load and process a random selection of test images
test_set = _load_data_set(config['data_path'], test=True)
selection = np.random.choice(test_set['rgb'].shape[0], size=n_imgs)
rgb = test_set['rgb'][selection]
depth = _preprocess_depth(test_set)[selection]
# Encode/reconstruct images and compute errors
reconstruction = model.predict(depth)
error = np.abs(depth - reconstruction)
# Plot results
fig = plt.figure()
grid = ImageGrid(fig, 111,
nrows_ncols=(n_imgs, 4),
share_all=True,
axes_pad=0.05,
cbar_mode='single',
cbar_location='right',
cbar_size='5%',
cbar_pad=None)
def _plot_sample(i, rgb, depth, reconstruction, error):
# Plot RGB
ax = grid[4 * i]
ax.set_axis_off()
ax.imshow(rgb)
def _add_depth_img(depth_img, j):
ax = grid[4 * i + j]
ax.set_axis_off()
img = ax.imshow(depth_img.squeeze(), cmap='viridis')
img.set_clim(0., 0.3)
ax.cax.colorbar(img)
# Plot depth, reconstruction and error
_add_depth_img(depth, 1)
_add_depth_img(reconstruction, 2)
_add_depth_img(error, 3)
for i in range(n_imgs):
_plot_sample(i, rgb[i], depth[i], reconstruction[i], error[i])
plt.savefig(os.path.join(args.model_dir, 'reconstructions.png'), dpi=300)
plt.show()
def test(args):
# Load the model
config = io_utils.load_yaml(os.path.join(args.model_dir, 'config.yaml'))
model = encoders.SimpleAutoEncoder(config)
model.load_weights(args.model_dir)
# Load the test set
test_set = _load_data_set(config['data_path'], test=True)
test_imgs = _preprocess_depth(test_set)
# Compute the test loss
loss = model.test(test_imgs, test_imgs)
print('Test loss: {}'.format(loss))
return loss
def main(args):
config = io_utils.load_yaml(args.config)
config['time_horizon'] = 150
config['simulation']['visualize'] = True
robot = RobotEnv(config)
if args.slider:
agent = slider_agent.SliderAgent(robot.action_space)
else:
agent = random_agent.RandomAgent(robot.action_space)
# Run and time the agent
run_agent(robot, agent, debug=True)
robot.close()
def from_yaml(file_path):
"""Read a transform from a yaml file.
Example of the content of such a file:
transform:
translation: [1., 2., 3.]
rotation: [0., 0., 0., 1.]
Args:
file_path: The path to the YAML file.
Returns:
A 4x4 homogeneous transformation matrix.
"""
cfg = io_utils.load_yaml(file_path)
return from_dict(cfg['transform'])
def run(args):
top_folder_idx = args.model.rfind('/')
top_folder_str = args.model[0:top_folder_idx]
config_file = top_folder_str + '/config.yaml'
config = io_utils.load_yaml(config_file)
normalize = config.get("normalize", False)
if args.visualize:
config['simulation']['real_time'] = False
config['simulation']['visualize'] = True
task = DummyVecEnv([
lambda: gym.make(
'gripper-env-v0', config=config, evaluate=True, test=args.test)
])
if normalize:
task = VecNormalize(task,
training=False,
norm_obs=False,
norm_reward=True,
clip_obs=10.)
task = VecNormalize.load(
os.path.join(top_folder_str, 'vecnormalize.pkl'), task)
# task = gym.make('gripper-env-v0', config=config, evaluate=True, test=args.test)
model_lower = args.model.lower()
if 'trpo' == config["algorithm"]:
agent = sb.TRPO.load(args.model)
elif 'sac' == config["algorithm"]:
agent = sb.SAC.load(args.model)
elif 'ppo' == config["algorithm"]:
agent = sb.PPO2.load(args.model)
elif 'dqn' == config["algorithm"]:
agent = sb.DQN.load(args.model)
elif 'bdq' == config["algorithm"]:
agent = sb.BDQ.load(args.model)
else:
raise Exception
print("Run the agent")
run_agent(task, agent, args.stochastic)
task.close()
def train(args):
# Load the encoder configuration
config = io_utils.load_yaml(args.config)
# If not existing, create the model directory
model_dir = os.path.expanduser(args.model_dir)
os.makedirs(model_dir, exist_ok=True)
# Build the model
model = encoders.SimpleAutoEncoder(config)
io_utils.save_yaml(config, os.path.join(model_dir, 'config.yaml'))
# Load and process the training data
train_set = _load_data_set(config['data_path'], test=False)
train_imgs = _preprocess_depth(train_set)
# Train the model
batch_size = config['batch_size']
epochs = config['epochs']
model.train(train_imgs, train_imgs, batch_size, epochs, model_dir)
def train(args):
config = io_utils.load_yaml(args.config)
os.mkdir(args.model_dir)
# Folder for best models
os.mkdir(args.model_dir + "/best_model")
model_dir = args.model_dir
algo = args.algo
if args.visualize:
config['simulation']['real_time'] = False
config['simulation']['visualize'] = True
if args.simple:
logging.info("Simplified environment is set")
config['simplified'] = True
if args.shaped:
logging.info("Shaped reward function is being used")
config['reward']['shaped'] = True
if args.timestep:
config[algo]['total_timesteps'] = args.timestep
if not args.algo == 'DQN':
config['robot']['discrete'] = False
else:
config['robot']['discrete'] = True
config[algo]['save_dir'] = model_dir
if args.timefeature:
env = DummyVecEnv([
lambda: TimeFeatureWrapper(
gym.make('gripper-env-v0', config=config))
])
else:
env = DummyVecEnv([
lambda: Monitor(gym.make('gripper-env-v0', config=config),
os.path.join(model_dir, "log_file"))
])
config["algorithm"] = args.algo.lower()
config_eval = config
config_eval['simulation']['real_time'] = False
config_eval['simulation']['visualize'] = False
io_utils.save_yaml(config, os.path.join(args.model_dir, 'config.yaml'))
io_utils.save_yaml(config,
os.path.join(args.model_dir, 'best_model/config.yaml'))
if args.timefeature:
test_env = DummyVecEnv([
lambda: TimeFeatureWrapper(
gym.make('gripper-env-v0',
config=config_eval,
evaluate=True,
validate=True))
])
else:
test_env = DummyVecEnv([
lambda: gym.make('gripper-env-v0',
config=config_eval,
evaluate=True,
validate=True)
])
sb_help = sb_helper.SBPolicy(env, test_env, config, model_dir,
args.load_dir, algo)
sb_help.learn()
env.close()
test_env.close()
