Compare commits
No commits in common. "c4283ced6fa9699ac94fe598a6dac935efc0d7f4" and "ada3cadf0c1e1242114befb688db8e0d5dc0d4b9" have entirely different histories.
c4283ced6f
...
ada3cadf0c
@ -194,27 +194,6 @@ class TransitionModel(nn.Module):
|
||||
prior = {"mean": state_prior_mean, "std": state_prior_std, "sample": sample_state_prior, "history": history, "distribution": state_prior_dist}
|
||||
return prior
|
||||
|
||||
def stack_states(self, states, dim=0):
|
||||
|
||||
s = dict(
|
||||
mean = torch.stack([state['mean'] for state in states], dim=dim),
|
||||
std = torch.stack([state['std'] for state in states], dim=dim),
|
||||
sample = torch.stack([state['sample'] for state in states], dim=dim),
|
||||
history = torch.stack([state['history'] for state in states], dim=dim),)
|
||||
dist = dict(distribution = [state['distribution'] for state in states])
|
||||
s.update(dist)
|
||||
return s
|
||||
|
||||
def imagine_rollout(self, state, action, history, horizon):
|
||||
imagined_priors = []
|
||||
for i in range(horizon):
|
||||
prior = self.imagine_step(state, action, history)
|
||||
state = prior["sample"]
|
||||
history = prior["history"]
|
||||
imagined_priors.append(prior)
|
||||
imagined_priors = self.stack_states(imagined_priors, dim=0)
|
||||
return imagined_priors
|
||||
|
||||
def reparemeterize(self, mean, std):
|
||||
eps = torch.randn_like(std)
|
||||
return mean + eps * std
|
||||
@ -248,6 +227,40 @@ class TanhBijector(torch.distributions.Transform):
|
||||
return 2.0 * (torch.log(torch.tensor([2.0])) - x - F.softplus(-2.0 * x))
|
||||
|
||||
|
||||
class CLUBSample(nn.Module): # Sampled version of the CLUB estimator
|
||||
def __init__(self, x_dim, y_dim, hidden_size):
|
||||
super(CLUBSample, self).__init__()
|
||||
self.p_mu = nn.Sequential(
|
||||
nn.Linear(x_dim, hidden_size//2),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_size//2, hidden_size//2),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_size//2, y_dim)
|
||||
)
|
||||
|
||||
self.p_logvar = nn.Sequential(
|
||||
nn.Linear(x_dim, hidden_size//2),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_size//2, hidden_size//2),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_size//2, y_dim),
|
||||
nn.Tanh()
|
||||
)
|
||||
|
||||
def get_mu_logvar(self, x_samples):
|
||||
mu = self.p_mu(x_samples)
|
||||
logvar = self.p_logvar(x_samples)
|
||||
return mu, logvar
|
||||
|
||||
def loglikeli(self, x_samples, y_samples):
|
||||
mu, logvar = self.get_mu_logvar(x_samples)
|
||||
return (-(mu - y_samples)**2 /logvar.exp()-logvar).sum(dim=1).mean(dim=0)
|
||||
|
||||
def forward(self, x_samples, y_samples):
|
||||
mu, logvar = self.get_mu_logvar(x_samples)
|
||||
return - self.loglikeli(x_samples, y_samples)
|
||||
|
||||
|
||||
class ProjectionHead(nn.Module):
|
||||
def __init__(self, state_size, action_size, hidden_size):
|
||||
super(ProjectionHead, self).__init__()
|
||||
@ -282,43 +295,3 @@ class ContrastiveHead(nn.Module):
|
||||
logits = logits - torch.max(logits, 1)[0][:, None]
|
||||
logits = logits * self.temperature
|
||||
return logits
|
||||
|
||||
|
||||
class CLUBSample(nn.Module): # Sampled version of the CLUB estimator
|
||||
def __init__(self, last_states, current_states, negative_current_states, predicted_current_states):
|
||||
super(CLUBSample, self).__init__()
|
||||
self.last_states = last_states
|
||||
self.current_states = current_states
|
||||
self.negative_current_states = negative_current_states
|
||||
self.predicted_current_states = predicted_current_states
|
||||
|
||||
def get_mu_var_samples(self, state_dict):
|
||||
dist = state_dict["distribution"]
|
||||
sample = dist.sample() # Use state_dict["sample"] if you want to use the same sample for all the losses
|
||||
mu = dist.mean
|
||||
var = dist.variance
|
||||
return mu, var, sample
|
||||
|
||||
def loglikeli(self):
|
||||
_, _, pred_sample = self.get_mu_var_samples(self.predicted_current_states)
|
||||
mu_curr, var_curr, _ = self.get_mu_var_samples(self.current_states)
|
||||
logvar_curr = torch.log(var_curr)
|
||||
return (-(mu_curr - pred_sample)**2 /var_curr-logvar_curr).sum(dim=1).mean(dim=0)
|
||||
|
||||
def forward(self):
|
||||
_, _, pred_sample = self.get_mu_var_samples(self.predicted_current_states)
|
||||
mu_curr, var_curr, _ = self.get_mu_var_samples(self.current_states)
|
||||
mu_neg, var_neg, _ = self.get_mu_var_samples(self.negative_current_states)
|
||||
|
||||
pos = (-(mu_curr - pred_sample)**2 /var_curr).sum(dim=1).mean(dim=0)
|
||||
neg = (-(mu_neg - pred_sample)**2 /var_neg).sum(dim=1).mean(dim=0)
|
||||
upper_bound = pos - neg
|
||||
|
||||
return upper_bound/2
|
||||
|
||||
def learning_loss(self):
|
||||
return - self.loglikeli()
|
||||
|
||||
|
||||
if "__name__ == __main__":
|
||||
pass
|
||||
140
DPI/train.py
140
DPI/train.py
@ -16,8 +16,6 @@ from logger import Logger
|
||||
from video import VideoRecorder
|
||||
from dmc2gym.wrappers import set_global_var
|
||||
|
||||
import torchvision.transforms as T
|
||||
|
||||
#from agent.baseline_agent import BaselineAgent
|
||||
#from agent.bisim_agent import BisimAgent
|
||||
#from agent.deepmdp_agent import DeepMDPAgent
|
||||
@ -33,7 +31,7 @@ def parse_args():
|
||||
parser.add_argument('--image_size', default=84, type=int)
|
||||
parser.add_argument('--channels', default=3, type=int)
|
||||
parser.add_argument('--action_repeat', default=1, type=int)
|
||||
parser.add_argument('--frame_stack', default=3, type=int)
|
||||
parser.add_argument('--frame_stack', default=4, type=int)
|
||||
parser.add_argument('--resource_files', type=str)
|
||||
parser.add_argument('--eval_resource_files', type=str)
|
||||
parser.add_argument('--img_source', default=None, type=str, choices=['color', 'noise', 'images', 'video', 'none'])
|
||||
@ -41,18 +39,18 @@ def parse_args():
|
||||
parser.add_argument('--high_noise', action='store_true')
|
||||
# replay buffer
|
||||
parser.add_argument('--replay_buffer_capacity', default=50000, type=int) #50000
|
||||
parser.add_argument('--episode_length', default=51, type=int)
|
||||
parser.add_argument('--episode_length', default=50, type=int)
|
||||
# train
|
||||
parser.add_argument('--agent', default='dpi', type=str, choices=['baseline', 'bisim', 'deepmdp', 'db', 'dpi', 'rpc'])
|
||||
parser.add_argument('--init_steps', default=10000, type=int)
|
||||
parser.add_argument('--num_train_steps', default=10000, type=int)
|
||||
parser.add_argument('--batch_size', default=20, type=int) #512
|
||||
parser.add_argument('--init_steps', default=1000, type=int)
|
||||
parser.add_argument('--num_train_steps', default=1000, type=int)
|
||||
parser.add_argument('--batch_size', default=200, type=int) #512
|
||||
parser.add_argument('--state_size', default=256, type=int)
|
||||
parser.add_argument('--hidden_size', default=128, type=int)
|
||||
parser.add_argument('--history_size', default=128, type=int)
|
||||
parser.add_argument('--num-units', type=int, default=200, help='num hidden units for reward/value/discount models')
|
||||
parser.add_argument('--load_encoder', default=None, type=str)
|
||||
parser.add_argument('--imagine_horizon', default=15, type=str)
|
||||
parser.add_argument('--imagination_horizon', default=15, type=str)
|
||||
# eval
|
||||
parser.add_argument('--eval_freq', default=10, type=int) # TODO: master had 10000
|
||||
parser.add_argument('--num_eval_episodes', default=20, type=int)
|
||||
@ -115,7 +113,6 @@ class DPI:
|
||||
|
||||
# environment setup
|
||||
self.env = make_env(self.args)
|
||||
#self.args.seed = np.random.randint(0, 1000)
|
||||
self.env.seed(self.args.seed)
|
||||
|
||||
# noiseless environment setup
|
||||
@ -193,123 +190,87 @@ class DPI:
|
||||
|
||||
def collect_sequences(self, episodes):
|
||||
obs = self.env.reset()
|
||||
#obs_clean = self.env_clean.reset()
|
||||
obs_clean = self.env_clean.reset()
|
||||
done = False
|
||||
|
||||
#video = VideoRecorder(self.video_dir if args.save_video else None, resource_files=args.resource_files)
|
||||
for episode_count in tqdm.tqdm(range(episodes), desc='Collecting episodes'):
|
||||
if args.save_video:
|
||||
self.env.video.init(enabled=True)
|
||||
#self.env_clean.video.init(enabled=True)
|
||||
self.env_clean.video.init(enabled=True)
|
||||
|
||||
for i in range(self.args.episode_length):
|
||||
|
||||
action = self.env.action_space.sample()
|
||||
|
||||
next_obs, rew, done, _ = self.env.step(action)
|
||||
#next_obs_clean, _, done, _ = self.env_clean.step(action)
|
||||
next_obs, _, done, _ = self.env.step(action)
|
||||
next_obs_clean, _, done, _ = self.env_clean.step(action)
|
||||
|
||||
self.data_buffer.add(obs, action, next_obs, episode_count+1, done)
|
||||
#self.data_buffer_clean.add(obs_clean, action, next_obs_clean, episode_count+1, done)
|
||||
self.data_buffer_clean.add(obs_clean, action, next_obs_clean, episode_count+1, done)
|
||||
|
||||
if args.save_video:
|
||||
self.env.video.record(self.env)
|
||||
#self.env_clean.video.record(self.env_clean)
|
||||
self.env.video.record(self.env_clean)
|
||||
self.env_clean.video.record(self.env_clean)
|
||||
|
||||
if done or i == self.args.episode_length-1:
|
||||
if done:
|
||||
obs = self.env.reset()
|
||||
#obs_clean = self.env_clean.reset()
|
||||
obs_clean = self.env_clean.reset()
|
||||
done=False
|
||||
else:
|
||||
obs = next_obs
|
||||
#obs_clean = next_obs_clean
|
||||
obs_clean = next_obs_clean
|
||||
if args.save_video:
|
||||
self.env.video.save('noisy/%d.mp4' % episode_count)
|
||||
#self.env_clean.video.save('clean/%d.mp4' % episode_count)
|
||||
self.env_clean.video.save('clean/%d.mp4' % episode_count)
|
||||
print("Collected {} random episodes".format(episode_count+1))
|
||||
|
||||
def train(self):
|
||||
# collect experience
|
||||
self.collect_sequences(self.args.batch_size)
|
||||
|
||||
# Group observations and next_observations by steps from past to present
|
||||
last_observations = torch.Tensor(self.data_buffer.group_steps(self.data_buffer,"observations")).float()[:self.args.episode_length-1]
|
||||
current_observations = torch.Tensor(self.data_buffer.group_steps(self.data_buffer,"next_observations")).float()[:self.args.episode_length-1]
|
||||
next_observations = torch.Tensor(self.data_buffer.group_steps(self.data_buffer,"next_observations")).float()[1:]
|
||||
actions = torch.Tensor(self.data_buffer.group_steps(self.data_buffer,"actions",obs=False)).float()[:self.args.episode_length-1]
|
||||
next_actions = torch.Tensor(self.data_buffer.group_steps(self.data_buffer,"actions",obs=False)).float()[1:]
|
||||
# Group observations and next_observations by steps
|
||||
observations = torch.Tensor(self.data_buffer.group_steps(self.data_buffer,"observations")).float()
|
||||
next_observations = torch.Tensor(self.data_buffer.group_steps(self.data_buffer,"next_observations")).float()
|
||||
actions = torch.Tensor(self.data_buffer.group_steps(self.data_buffer,"actions",obs=False)).float()
|
||||
|
||||
# Initialize transition model states
|
||||
self.transition_model.init_states(self.args.batch_size, device="cpu") # (N,128)
|
||||
self.history = self.transition_model.prev_history # (N,128)
|
||||
|
||||
# Train encoder
|
||||
total_ub_loss = 0
|
||||
total_encoder_loss = 0
|
||||
for i in range(self.args.episode_length-1):
|
||||
if i > 0:
|
||||
previous_information_loss = 0
|
||||
previous_encoder_loss = 0
|
||||
for i in range(self.args.episode_length):
|
||||
# Encode observations and next_observations
|
||||
self.last_states_dict = self.obs_encoder(last_observations[i])
|
||||
self.current_states_dict = self.obs_encoder(current_observations[i])
|
||||
self.next_states_dict = self.obs_encoder_momentum(next_observations[i])
|
||||
self.action = actions[i] # (N,6)
|
||||
history = self.transition_model.prev_history
|
||||
self.states_dist = self.obs_encoder(observations[i])
|
||||
self.next_states_dist = self.obs_encoder(next_observations[i])
|
||||
|
||||
# Encode negative observations
|
||||
idx = torch.randperm(current_observations[i].shape[0]) # random permutation on batch
|
||||
random_time_index = torch.randint(0, self.args.episode_length-2, (1,)).item() # random time index
|
||||
negative_current_observations = current_observations[random_time_index][idx]
|
||||
self.negative_current_states_dict = self.obs_encoder(negative_current_observations)
|
||||
|
||||
# Predict current state from past state with transition model
|
||||
last_states_sample = self.last_states_dict["sample"]
|
||||
predicted_current_state_dict = self.transition_model.imagine_step(last_states_sample, self.action, self.history)
|
||||
self.history = predicted_current_state_dict["history"]
|
||||
# Sample states and next_states
|
||||
self.states = self.states_dist["sample"] # (N,128)
|
||||
self.next_states = self.next_states_dist["sample"] # (N,128)
|
||||
self.actions = actions[i] # (N,6)
|
||||
|
||||
# Calculate upper bound loss
|
||||
ub_loss = self._upper_bound_minimization(self.last_states_dict,
|
||||
self.current_states_dict,
|
||||
self.negative_current_states_dict,
|
||||
predicted_current_state_dict
|
||||
)
|
||||
past_latent_loss = previous_information_loss + self._upper_bound_minimization(self.states, self.next_states)
|
||||
|
||||
# Calculate encoder loss
|
||||
encoder_loss = self._past_encoder_loss(self.current_states_dict,
|
||||
predicted_current_state_dict)
|
||||
|
||||
total_ub_loss += ub_loss
|
||||
total_encoder_loss += encoder_loss
|
||||
|
||||
imagine_horizon = np.minimum(self.args.imagine_horizon, self.args.episode_length-1-i)
|
||||
imagined_rollout = self.transition_model.imagine_rollout(self.current_states_dict["sample"], self.action, self.history, imagine_horizon)
|
||||
print(imagine_horizon)
|
||||
#exit()
|
||||
|
||||
#print(total_ub_loss, total_encoder_loss)
|
||||
past_encoder_loss = previous_encoder_loss + self._past_encoder_loss(self.states, self.next_states,
|
||||
self.states_dist, self.next_states_dist,
|
||||
self.actions, self.history, i)
|
||||
|
||||
|
||||
print(past_encoder_loss, past_latent_loss)
|
||||
|
||||
def _upper_bound_minimization(self, last_states, current_states, negative_current_states, predicted_current_states):
|
||||
club_sample = CLUBSample(last_states,
|
||||
current_states,
|
||||
negative_current_states,
|
||||
predicted_current_states)
|
||||
club_loss = club_sample()
|
||||
previous_information_loss = past_latent_loss
|
||||
previous_encoder_loss = past_encoder_loss
|
||||
|
||||
def _upper_bound_minimization(self, states, next_states):
|
||||
club_sample = CLUBSample(self.args.state_size,
|
||||
self.args.state_size,
|
||||
self.args.hidden_size)
|
||||
club_loss = club_sample(states, next_states)
|
||||
return club_loss
|
||||
|
||||
def _past_encoder_loss(self, curr_states_dict, predicted_curr_states_dict):
|
||||
# current state distribution
|
||||
curr_states_dist = curr_states_dict["distribution"]
|
||||
|
||||
# predicted current state distribution
|
||||
predicted_curr_states_dist = predicted_curr_states_dict["distribution"]
|
||||
|
||||
# KL divergence loss
|
||||
loss = torch.distributions.kl.kl_divergence(curr_states_dist, predicted_curr_states_dist).mean()
|
||||
|
||||
return loss
|
||||
|
||||
"""
|
||||
def _past_encoder_loss(self, states, next_states, states_dist, next_states_dist, actions, history, step):
|
||||
# Imagine next state
|
||||
if step == 0:
|
||||
@ -326,21 +287,6 @@ class DPI:
|
||||
loss = torch.distributions.kl.kl_divergence(imagined_next_states_dist, next_states_dist["distribution"]).mean()
|
||||
|
||||
return loss
|
||||
"""
|
||||
|
||||
def get_features(self, x, momentum=False):
|
||||
if self.aug:
|
||||
x = T.RandomCrop((80, 80))(x) # (None,80,80,4)
|
||||
x = T.functional.pad(x, (4, 4, 4, 4), "symmetric") # (None,88,88,4)
|
||||
x = T.RandomCrop((84, 84))(x) # (None,84,84,4)
|
||||
|
||||
with torch.no_grad():
|
||||
x = (x.float() - self.ob_mean) / self.ob_std
|
||||
if momentum:
|
||||
x = self.obs_encoder(x).detach()
|
||||
else:
|
||||
x = self.obs_encoder_momentum(x)
|
||||
return x
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = parse_args()
|
||||
|
||||
@ -161,11 +161,11 @@ class ReplayBuffer:
|
||||
non_zero_indices = np.nonzero(buffer.episode_count)[0]
|
||||
variable = variable[non_zero_indices]
|
||||
if obs:
|
||||
variable = variable.reshape(self.args.batch_size, self.args.episode_length,
|
||||
variable = variable.reshape(self.args.episode_length, self.args.batch_size,
|
||||
self.args.frame_stack*self.args.channels,
|
||||
self.args.image_size,self.args.image_size).transpose(1, 0, 2, 3, 4)
|
||||
self.args.image_size,self.args.image_size)
|
||||
else:
|
||||
variable = variable.reshape(self.args.batch_size, self.args.episode_length, -1).transpose(1, 0, 2)
|
||||
variable = variable.reshape(self.args.episode_length, self.args.batch_size,-1)
|
||||
return variable
|
||||
|
||||
def transform_grouped_steps(self, variable):
|
||||
|
||||
Loading…
Reference in New Issue
Block a user