VedantDave/Curiosity
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Implementing ICLUB

Browse Source
This commit is contained in:
Vedant Dave 2023-03-24 20:39:14 +01:00
parent abaca2bea9
commit 641c9bd57c
3 changed files with 77 additions and 28 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
DPI/models.py
View File

@ -128,26 +128,28 @@ class TransitionModel(nn.Module):
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),
self.p_mu = nn.Sequential(
nn.Linear(x_dim, hidden_size//2),
nn.ReLU(),
nn.Linear(hidden_size//2, y_dim))
nn.Linear(hidden_size//2, y_dim)
)
self.p_logvar = nn.Sequential(nn.Linear(x_dim, hidden_size//2),
self.p_logvar = nn.Sequential(
nn.Linear(x_dim, hidden_size//2),
nn.ReLU(),
nn.Linear(hidden_size//2, y_dim),
nn.Tanh())
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)
@ -165,8 +167,9 @@ class CLUBSample(nn.Module): # Sampled version of the CLUB estimator
if __name__ == "__main__":
encoder = ObservationEncoder((12,84,84), 256)
x = torch.randn(100, 12, 84, 84)
x = torch.randn(5000, 12, 84, 84)
print(encoder(x).shape)
exit()
club = CLUBSample(256, 256 , 512)
x = torch.randn(100, 256)

50
DPI/train.py
View File

@ -9,7 +9,7 @@ import dmc2gym
import wandb
import utils
from utils import ReplayBuffer, make_env
from utils import ReplayBuffer, make_env, save_image
from models import ObservationEncoder, ObservationDecoder, TransitionModel, CLUBSample
from logger import Logger
from video import VideoRecorder
@ -34,18 +34,18 @@ def parse_args():
parser.add_argument('--img_source', default=None, type=str, choices=['color', 'noise', 'images', 'video', 'none'])
parser.add_argument('--total_frames', default=1000, type=int)
# replay buffer
parser.add_argument('--replay_buffer_capacity', default=100000, type=int)
parser.add_argument('--episode_length', default=1000, type=int)
parser.add_argument('--replay_buffer_capacity', default=50000, type=int) #100000
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=1000, type=int)
parser.add_argument('--num_train_steps', default=1000, type=int)
parser.add_argument('--batch_size', default=512, 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('--k', default=3, type=int, help='number of steps for inverse model')
parser.add_argument('--load_encoder', default=None, 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)
@ -79,7 +79,6 @@ def parse_args():
parser.add_argument('--alpha_beta', default=0.9, type=float)
# misc
parser.add_argument('--seed', default=1, type=int)
parser.add_argument('--seed_steps', default=5000, type=int)
parser.add_argument('--work_dir', default='.', type=str)
parser.add_argument('--save_tb', default=False, action='store_true')
parser.add_argument('--save_model', default=False, action='store_true')
@ -117,20 +116,21 @@ class DPI:
obs_shape=(self.args.frame_stack*self.args.channels,self.args.image_size,self.args.image_size),
action_size=self.env.action_space.shape[0],
seq_len=self.args.episode_length,
batch_size=args.batch_size)
batch_size=args.batch_size,
args=self.args)
# create work directory
utils.make_dir(self.args.work_dir)
video_dir = utils.make_dir(os.path.join(self.args.work_dir, 'video'))
model_dir = utils.make_dir(os.path.join(self.args.work_dir, 'model'))
buffer_dir = utils.make_dir(os.path.join(self.args.work_dir, 'buffer'))
self.video_dir = utils.make_dir(os.path.join(self.args.work_dir, 'video'))
self.model_dir = utils.make_dir(os.path.join(self.args.work_dir, 'model'))
self.buffer_dir = utils.make_dir(os.path.join(self.args.work_dir, 'buffer'))
# create video recorder
#video = VideoRecorder(video_dir if args.save_video else None, resource_files=args.resource_files)
#video.init(enabled=True)
# create models
self.build_models(use_saved=False, saved_model_dir=model_dir)
self.build_models(use_saved=False, saved_model_dir=self.model_dir)
def build_models(self, use_saved, saved_model_dir=None):
self.obs_encoder = ObservationEncoder(
@ -174,6 +174,7 @@ class DPI:
next_obs, _, done, _ = self.env.step(action)
self.data_buffer.add(obs, action, next_obs, episode_count+1, done)
if done:
obs = self.env.reset()
done=False
@ -185,12 +186,33 @@ class DPI:
#video.save('%d.mp4' % step)
#video.close()
def upper_bound_minimization(self):
pass
def train(self):
# collect experience
self.collect_random_episodes(self.args.batch_size)
# 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()
# Train encoder
for i in range(self.args.episode_length):
# Encode observations and next_observations
self.features = self.obs_encoder(observations[i]) # (N,128)
self.next_features = self.obs_encoder(next_observations[i]) # (N,128)
# Calculate upper bound loss
past_loss = self.upper_bound_minimization(self.features, self.next_features)
def upper_bound_minimization(self, features, next_features):
club_sample = CLUBSample(self.args.state_size,
self.args.state_size,
self.args.hidden_size)
club_loss = club_sample(features, next_features)
return club_loss
if __name__ == '__main__':
args = parse_args()
dpi = DPI(args)
dpi.collect_random_episodes(episodes=5)
dpi.train()

26
DPI/utils.py
View File

@ -13,6 +13,7 @@ import gym
import dmc2gym
import random
from PIL import Image
from collections import deque
@ -105,7 +106,7 @@ class FrameStack(gym.Wrapper):
class ReplayBuffer:
def __init__(self, size, obs_shape, action_size, seq_len, batch_size):
def __init__(self, size, obs_shape, action_size, seq_len, batch_size, args):
self.size = size
self.obs_shape = obs_shape
self.action_size = action_size
@ -113,6 +114,7 @@ class ReplayBuffer:
self.batch_size = batch_size
self.idx = 0
self.full = False
self.args = args
self.observations = np.empty((size, *obs_shape), dtype=np.uint8)
self.actions = np.empty((size, action_size), dtype=np.float32)
self.next_observations = np.empty((size, *obs_shape), dtype=np.uint8)
@ -152,6 +154,22 @@ class ReplayBuffer:
obs,acs,rews,terms= self._retrieve_batch(np.asarray([self._sample_idx(l) for _ in range(n)]), n, l)
return obs,acs,rews,terms
def group_steps(self, buffer, variable):
variable = getattr(buffer, variable)
non_zero_indices = np.nonzero(buffer.episode_count)[0]
variable = variable[non_zero_indices]
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)
return variable
def transform_grouped_steps(self, variable):
variable = variable.transpose((1, 0, 2, 3, 4))
variable = variable.reshape(self.args.batch_size*self.args.episode_length,self.args.frame_stack*self.args.channels,
self.args.image_size,self.args.image_size)
return variable
def make_env(args):
env = dmc2gym.make(
@ -168,3 +186,9 @@ def make_env(args):
frame_skip=args.action_repeat
)
return env
def save_image(array, filename):
array = array.transpose(1, 2, 0)
array = (array * 255).astype(np.uint8)
image = Image.fromarray(array)
image.save(filename)