def show_reconstruction(cfg, args):
vae_path = args['--vae']
print('vae_path: %s' % vae_path)
if not vae_path:
print('Error: No vae path specified')
sys.exit(0)
# init vae
print('Initializing vae...')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
vae = VAE(image_channels=cfg.IMAGE_CHANNELS, z_dim=cfg.VARIANTS_SIZE)
vae.load_state_dict(torch.load(vae_path,
map_location=torch.device(device)))
vae.to(device).eval()
for i in range(cfg.TIME_STEPS):
z = torch.load('z_tensor.pt')
reconst = vae.decode(z)
reconst = reconst.detach().cpu()[0].numpy()
reconst = np.transpose(np.uint8(reconst * 255), [1, 2, 0])
reconst_image = F_.to_pil_image(reconst)
imgplot = plt.imshow(reconst_image)
plt.pause(0.05)
time.sleep(0.1)
def observe_and_learn(cfg, model_path, vae_path=None):
global ctr
time_steps = 5000
time_steps = 10000
try:
time.sleep(5)
vae = None
device = None
if vae_path:
# init vae
print('Initializing vae...')
device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
vae = VAE(image_channels=cfg.IMAGE_CHANNELS,
z_dim=cfg.VARIANTS_SIZE)
vae.load_state_dict(
torch.load(vae_path, map_location=torch.device(device)))
vae.to(device).eval()
# create agent; wrapper for environment; later we can add vae to the agent
agent = DonkeyAgent(cam,
time_step=0.05,
frame_skip=1,
env_type='simulator',
controller=ctr,
vae=vae,
device=device)
print('DonkeyAgent created...')
model = CustomSAC(CustomSACPolicy,
agent,
verbose=cfg.VERBOSE,
batch_size=cfg.BATCH_SIZE,
buffer_size=cfg.BUFFER_SIZE,
learning_starts=cfg.LEARNING_STARTS,
gradient_steps=cfg.GRADIENT_STEPS,
train_freq=cfg.TRAIN_FREQ,
ent_coef=cfg.ENT_COEF,
learning_rate=cfg.LEARNING_RATE)
print('CustomSAC Initialized.')
print('learning...')
ctr.mode = 'local'
model.learn(total_timesteps=cfg.TIME_STEPS,
log_interval=cfg.LOG_INTERVAL)
model.save(model_path)
print('Model finished.')
except Exception as e:
print('error:no new model generated. %s' % e)
traceback.print_exc()
class WorldModel(nn.Module):
def __init__(self,
action_dim: int,
output_dim: int,
vae_path: str = './vae/model/best.tar',
mdnrnn_path: str = './mdnrnn/model/best.tar',
activation: Callable = torch.relu,
output_activation: Union[Callable, str] = None,
output_squeeze: bool = False) -> None:
super(WorldModel, self).__init__()
# define VAE model
self.latent_size = LATENT_SIZE
self.vae = VAE(CHANNELS, LATENT_SIZE)
vae_state = torch.load(vae_path)
self.vae.load_state_dict(vae_state['state_dict'])
for param in self.vae.parameters():
param.requires_grad = False
# define MDNRNN model
self.n_gauss = N_GAUSS
self.mdnrnn = MDNRNN(LATENT_SIZE, action_dim, HIDDEN_SIZE, N_GAUSS, rewards_terminal=False)
mdnrnn_state = torch.load(mdnrnn_path)
self.mdnrnn.load_state_dict(mdnrnn_state['state_dict'])
for param in self.mdnrnn.parameters():
param.requires_grad = False
# controller
ctr_size = LATENT_SIZE + N_GAUSS + 2 * (LATENT_SIZE * N_GAUSS)
self.fc1 = nn.Linear(ctr_size, 512)
self.fc2 = nn.Linear(512, output_dim)
self.activation = activation
self.output_activation = output_activation
self.output_squeeze = output_squeeze
def forward(self, x: torch.Tensor, a: torch.Tensor) -> torch.Tensor:
# VAE
_, mu, logsigma = self.vae(x)
latent = (mu + logsigma.exp() * torch.randn_like(mu)).view(-1, self.latent_size)
# MDNRNN
mus, sigmas, logpi = self.mdnrnn(a.unsqueeze(0), latent.unsqueeze(0))
# reshape
mus = torch.squeeze(mus, dim=0).view(-1, self.n_gauss * self.latent_size)
sigmas = torch.squeeze(sigmas, dim=0).view(-1, self.n_gauss * self.latent_size)
logpi = torch.squeeze(logpi, dim=0).view(-1, self.n_gauss)
# controller
ctr_in = [latent, mus, sigmas, logpi]
x = torch.cat(ctr_in, dim=1)
x = self.activation(self.fc1(x))
if self.output_activation is None:
x = self.fc2(x)
else:
x = self.output_activation(self.fc2(x))
return x.squeeze() if self.output_squeeze else x
def main():
args = parse_args()
input_size = (settings.reduced_image_channels,
settings.reduced_image_width, settings.reduced_image_height)
vae = VAE(input_size=input_size,
latent_dim=settings.vae_latent_dim).to(settings.device)
savefile = Path(args.savefile)
if savefile.exists():
vae.load_state_dict(torch.load(f'{savefile}'))
vae.eval()
run(vae, savefile)
def optimize_model(cfg,
model_path,
vae_path=None,
auto_mode=0,
env_type='simulator'):
global ctr
time_steps = 5000
vae = None
if vae_path:
# init vae
print('Initializing vae...')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
vae = VAE(image_channels=cfg.IMAGE_CHANNELS, z_dim=cfg.VARIANTS_SIZE)
vae.load_state_dict(
torch.load(vae_path, map_location=torch.device(device)))
vae.to(device).eval()
# create agent; wrapper for environment; later we can add vae to the agent
agent = DonkeyAgent(cam,
time_step=0.05,
frame_skip=1,
env_type=env_type,
controller=ctr,
vae=vae,
auto_mode=auto_mode)
print('DonkeyAgent created...')
model = CustomSAC.load(model_path, agent)
print('Executing model...')
model.learning_starts = 0
ctr.mode = 'local'
model.learn(total_timesteps=time_steps, log_interval=cfg.LOG_INTERVAL)
model.save(model_path)
print('Model finished.')
def drive_model(cfg, model_path, vae_path=None, env_type='simulator'):
global ctr
time_steps = 10000
vae = None
if vae_path:
# init vae
print('Initializing vae...')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
vae = VAE(image_channels=cfg.IMAGE_CHANNELS, z_dim=cfg.VARIANTS_SIZE)
vae.load_state_dict(
torch.load(vae_path, map_location=torch.device(device)))
vae.to(device).eval()
# create agent; wrapper for environment; later we can add vae to the agent
agent = DonkeyAgent(cam,
time_step=0.05,
frame_skip=1,
env_type='simulator',
controller=ctr,
vae=vae)
print('DonkeyAgent created...')
model = CustomSAC.load(model_path)
print('Executing model...')
obs = agent.reset()
ctr.mode = 'local'
for step in range(time_steps):
if step % 100 == 0: print("step: ", step)
action, _states = model.predict(obs, deterministic=False)
obs, rewards, dones, info = agent.step(action)
time.sleep(0.05)
while agent.is_game_over():
print('waiting for control')
time.sleep(1)
class MyEnv:
def __init__(self, env_):
self.env = env_
self.dev = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
self.action_space = env_.action_space
self._state_steps = 3
self.observation_space = (80, 160, 3*self._state_steps)
print("obs shape {}".format(self.observation_space))
self.state_shape = 32*self._state_steps
# vae
self.vae = VAE()
model_path = "./vae/vae_ikeda_ver1_32.pth"
self.vae.load_state_dict(torch.load(model_path))
self.vae.to(self.dev)
self._gen_id = 0 # 何回目のgenerateかを保持
self._frames = [] # mp4生成用にframeを保存
# pre processing
self._state_frames = deque(maxlen=self._state_steps) # 変換前のframe
self._gen_id = 0 # 何回目のgenerateかを保持
self._frames = [] # mp4生成用にframeを保存
self._step_repeat_times = 2
# self.num_envs = 1
def close(self):
self.env.close()
def step(self, action, show=False):
rews = 0.0
for i in range(self._step_repeat_times):
n_state, rew, done, info = self.env.step(action)
rews += rew
if i == 0:
self._state_frames.append(self.adjust_picture(n_state))
if show:
self._frames.append(np.array(n_state))
if done:
break
n_state_return = self.convert_state() # state 生成
rew = self.change_rew(rews/self._step_repeat_times, info)
if info["cte"] > 2.5:
done = True
rew = -1.0
elif info["cte"] < -5.0:
done = True
rew = -1.0
return n_state_return, rew, done, info
def change_rew(self, rew, info):
if info["speed"] < 0.0:
return -0.6
if abs(info["cte"]) < 1.0:
rew = info["speed"] / 100.0
else:
rew = -0.6
return rew
def reset(self):
rand_step = random.randrange(10)
self.env.reset()
for _ in range(rand_step + self._state_steps):
action = self.env.action_space.sample()
for i in range(self._step_repeat_times):
n_state, _, _, _ = self.env.step(action)
if i == 0:
self._state_frames.append(self.adjust_picture(n_state))
state = self.convert_state()
return state
def seed(self, seed_):
self.env.seed(seed_)
def adjust_picture(self, pict):
vae_state = self.convert_state_vae(pict)
return vae_state
def convert_state(self):
state_pre = []
for state in self._state_frames:
state_pre.append(state)
state = np.concatenate(state_pre, 0)
return state
def convert_state_to_tensor(self, state): # state(array) -> np.array -> convert some -> tensor
state_ = np.array(state).reshape((160, 120, 3))
state_ = state_[0:160, 40:120, :].reshape((1, 80, 160, 3))
state_ = torch.from_numpy(state_).permute(0, 3, 1, 2).float().to(self.dev)
return state_
def convert_state_vae(self, state):
state_ = self.convert_state_to_tensor(state)
state_, _, _ = self.vae.encode(state_)
state_ = state_.clone().detach().cpu().numpy()[0]
return state_
def generate_mp4(self):
# for mp4
fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
fps = 10
video = cv2.VideoWriter('./mp4/output' + str(self._gen_id) + ".mp4", fourcc, fps, (120, 160))
if not video.isOpened():
print("can't be opened")
# for path
os.mkdir("./tmp")
current_path = os.getcwd() # 現在のディレクトリ
# main procedure
for idx, frame in enumerate(self._frames):
fram = Image.fromarray(frame, "RGB")
path = current_path + "/tmp/frame" + str(idx) + ".png"
fram.save(path, 'PNG')
img = cv2.imread(path)
img = cv2.resize(img, (120, 160))
if img is None:
print("can't read")
break
video.write(img)
video.release()
shutil.rmtree("./tmp")
self._frames.clear()
self._gen_id += 1
def _load_vae(model_path, variants_size, image_channels, device):
vae = VAE(image_channels=image_channels, z_dim=variants_size)
vae.load_state_dict(
torch.load(model_path, map_location=torch.device(device)))
vae.to(torch.device(device)).eval()
return vae
class MyEnv:
def __init__(self, env_):
self.env = env_
self.dev = torch.device(
'cuda:0' if torch.cuda.is_available() else 'cpu')
self.action_space = env_.action_space
# self.observation_space = (1, 32)
# vae
self.vae = VAE()
model_path = "./vae/vae.pth"
self.vae.load_state_dict(torch.load(model_path))
self.vae.to(self.dev)
self._gen_id = 0 # 何回目のgenerateかを保持
self._frames = [] # mp4生成用にframeを保存
def step(self, action, show=False):
n_state, rew, done, info = self.env.step(action)
if show:
self._frames.append(np.array(n_state))
n_state = self.convert_state_vae(n_state)
rew = self.change_rew(rew, info)
if info["cte"] > 3.5:
done = True
rew = -1.0
elif info["cte"] < -5.0:
done = True
rew = -1.0
return n_state, rew, done, info
def change_rew(self, rew, info):
if info["speed"] < 0.0:
return -0.6
elif abs(info["cte"]) >= 2.0:
return -1.0
if rew > 0.0:
rew /= 20.0
if info["speed"] > 3.0:
rew += info["speed"] / 30.0
return rew
def reset(self):
state = self.env.reset()
state = self.convert_state_vae(state)
return state
def seed(self, seed_):
self.env.seed(seed_)
def convert_state_to_tensor(
self, state): # state(array) -> np.array -> convert some -> tensor
state_ = np.array(state).reshape((160, 120, 3))
# print("state_ shape {}".format(state1.shape))
state_ = state_[0:160, 40:120, :].reshape((1, 80, 160, 3))
# print("state shape {}".format(state_.shape))
# state_ = state_.reshape((1, 80, 160, 3))
state_ = torch.from_numpy(state_).permute(0, 3, 1,
2).float().to(self.dev)
state_ /= 255.0
return state_
def convert_state_vae(self, state):
state_ = self.convert_state_to_tensor(state)
state_, _, _ = self.vae.encode(state_)
state_ = state_.clone().detach().cpu().numpy()[0]
return state_
def generate_mp4(self):
# for mp4
fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
fps = 10
current_path = os.getcwd() # 現在のディレクトリ
video = cv2.VideoWriter(
current_path + '/mp4/output' + str(self._gen_id) + ".mp4", fourcc,
fps, (120, 160))
if not video.isOpened():
print("can't be opened")
# for path
os.mkdir("./tmp")
current_path = os.getcwd() # 現在のディレクトリ
# main procedure
for idx, frame in enumerate(self._frames):
fram = Image.fromarray(frame, "RGB")
path = current_path + "/tmp/frame" + str(idx) + ".png"
fram.save(path, 'PNG')
img = cv2.imread(path)
img = cv2.resize(img, (120, 160))
if img is None:
print("can't read")
break
video.write(img)
video.release()
shutil.rmtree("./tmp")
self._frames.clear()
self._gen_id += 1