def main(dataset_size, action_encoder_dir, output_dir):
outputdir = Path(output_dir)
outputdir.mkdir(parents=True, exist_ok=True)
screen_viewer = ImageViewer('input', (320, 480), format='tensor_gym_RGB')
dataset = ActionEncoderDataset(action_encoder_dir)
loader = data_utils.DataLoader(dataset=dataset,
batch_size=1,
shuffle=False,
drop_last=True,
collate_fn=collate_action_observation)
num_images = 0
for episode, (screen, observation, action, reward, done,
latent) in enumerate(loader):
if num_images >= dataset_size:
break
step = 10
length = screen[0].shape[0]
length = length - length % step
index = torch.Tensor(range(0, length, step)).long()
offsets = torch.randint_like(index, 0, step - 1)
index = index + offsets
sampled_screen = screen[0][index]
sampled_reward = reward[0][index]
for i, (frame, reward) in enumerate(zip(sampled_screen,
sampled_reward)):
if num_images >= dataset_size:
break
w_frame = NumpyRGBWrapper(frame, 'tensor_gym_RGB')
file = outputdir / Path('pic%04d_%04d' %
(episode, i)).with_suffix('.png')
imageio.imwrite(file, w_frame.numpyRGB)
file = outputdir / Path('rew%04d_%04d' %
(episode, i)).with_suffix('.np')
with file.open(mode='wb') as f:
pickle.dump(reward, f)
screen_viewer.update(frame)
num_images += 1
print(num_images)
def test_rollout_gen(self):
env = gym.make('SpaceInvaders-v4')
models = config.basepath() / 'SpaceInvaders-v4' / 'models'
visualsfile = models / 'GM53H301W5YS38XH'
visuals = Storeable.load(str(visualsfile)).to(config.device())
controllerfile = models / 'best_model68'
controller = torch.load(str(controllerfile))
policy = VCPolicy(visuals, controller)
viewer = ImageViewer('screen', (420, 360), 'numpyRGB')
for screen, observation, reward, done, info, action in RolloutGen(
env, policy):
viewer.update(screen)
def main(gym_environment, policy, output_dir):
global action_encoder, input_viewer
env = gym.make(gym_environment)
action_encoder = ActionEncoder(env, gym_environment, ActionEmbedding(env), policy.v).to(config.device())
rollout = Rollout(env)
input_viewer = ImageViewer('input', (320, 480), 'numpyRGB')
def frame(step):
input_viewer.update(step.screen)
rollout.register_before_hook(frame)
def save(args):
action_encoder.save_session()
def action_encoder_frame(step):
episode = step.meta['episode']
file_path = config.basepath() / gym_environment / output_dir / str(episode)
step.meta['filename'] = str(file_path)
action_encoder.update(step)
rollout.register_step(action_encoder_frame)
rollout.register_end_session(save)
for i_episode in range(531, 1000):
rollout.rollout(policy, max_timesteps=3000, episode=i_episode)
from mentalitystorm.config import config
from mentalitystorm.observe import ImageViewer
from mentalitystorm.data import DataPackage
from mentalitystorm.runners import SimpleRunFac, Params, Run
from mentalitystorm.util import Handles
import mentalitystorm.transforms as tf
from mentalitystorm.losses import MSELoss
import torchvision.transforms as TVT
from models import Compressor
from tqdm import tqdm
from torch.optim import Adam
if __name__ == '__main__':
input_viewer = ImageViewer('input', (320, 480))
output_viewer = ImageViewer('output', (320, 480))
latent_viewer = ImageViewer('latent', (320, 480))
latent_instr = LatentInstrument()
invader = tf.ColorMask(lower=[120, 125, 30],
upper=[140, 140, 130],
append=False)
cut_invader = tf.SetRange(0, 30, 0, 210, [0])
co_ord_conv_shots = GymImageDataset(
directory=config.datapath(r'SpaceInvaders-v4\images\raw_v1\all'),
input_transform=TVT.Compose(
[invader, cut_invader,
TVT.ToTensor(),
tf.CoordConv()]),
from mentalitystorm.observe import ImageViewer
view_latent1 = ImageViewer('latent1', (320, 480), channels=[0, 1, 2])
view_latent2 = ImageViewer('latent2', (320, 480), channels=[3])
view_input1 = ImageViewer('input1', (320, 480), channels=[0, 1, 2])
view_input2 = ImageViewer('input2', (320, 480), channels=[3, 4, 5])
def view_image(model, input, output):
view_input1.update(input[0].data)
view_input2.update(input[0].data)
view_latent1.update(output[0].data)
view_latent2.update(output[0].data)
decode_viewer1 = ImageViewer('decoded1', (320, 480), channels=[0, 1, 2])
decode_viewer2 = ImageViewer('decoded2', (320, 480), channels=[3])
def view_decode(model, input, output):
image = model.decode(output)
decode_viewer1.update(image)
decode_viewer2.update(image)
from mentalitystorm.instrumentation import register_tb, LatentInstrument
from mentalitystorm.data import StandardSelect, GymImageDataset, DataPackage
from mentalitystorm.config import config
from mentalitystorm.observe import ImageViewer
from mentalitystorm.runners import Run, SimpleRunFac, Params
import mentalitystorm.transforms as tf
import torchvision.transforms as TVT
from tqdm import tqdm
from torch.optim import Adam
from mentalitystorm.basemodels import MultiChannelAE
from mentalitystorm.train import SimpleInference
if __name__ == '__main__':
input_viewer = ImageViewer('input', (320, 480), channels=[0, 1, 2])
input2_viewer = ImageViewer('input2', (320, 480), channels=[3, 4, 5])
shot_viewer = ImageViewer('shot', (320, 480), channels=[0])
player_viewer = ImageViewer('player', (320, 480), channels=[1])
invader_viewer = ImageViewer('invader', (320, 480), channels=[2])
barrier_viewer = ImageViewer('barrier', (320, 480), channels=[3])
decode_viewer = ImageViewer('decoded', (320, 480), channels=[0, 1, 2])
decode2_viewer = ImageViewer('decoded2', (320, 480), channels=[3])
latent_instr = LatentInstrument()
def view_decode(model, input, output):
image = model.decode(output)
decode_viewer.update(image)
decode2_viewer.update(image)
shots = tf.ColorMask(lower=[128, 128, 128], upper=[255, 255, 255], append=True)
from mentalitystorm.data import ActionEncoderDataset, collate_action_observation
import torch.utils.data as data_utils
from mentalitystorm.observe import ImageViewer
screen_viewer = ImageViewer('input', (320, 480), 'tensor_gym_RGB')
dataset = ActionEncoderDataset(r'c:\data\SpaceInvaders-v4\rl_raw_v1')
loader = data_utils.DataLoader(dataset=dataset, batch_size=1, shuffle=False, drop_last=True,
collate_fn=collate_action_observation)
for screen, observation, action, reward, done, latent in loader:
for frame in screen[0]:
screen_viewer.update(frame)