wassname/iris_bigvae
1
0
Fork 0
mirror of https://github.com/wassname/iris_bigvae.git synced 2026-09-09 11:24:31 +08:00
Code Issues Packages Projects Releases Wiki Activity

checkin

Browse Source
This commit is contained in:
wassname
2024-05-11 10:11:10 +08:00
parent 73fdfddacb
commit a4234425af
13 changed files with 2345 additions and 1515 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+2
View File
@@ -134,3 +134,5 @@ dmypy.json
# Pyre type checker
.pyre/
bigvae_transformer_wm_research_log.md
journal.md
.vscode/launch.json
Vendored
+4 -3
View File
@@ -12,11 +12,11 @@
"console": "integratedTerminal",
"justMyCode": false,
"autoReload": {"enable": true,},
"env": {"WANDB_MODE":"disabled"},
"env": {"WANDB_MODE":"disabled","WANDB_SILENT":"true","HYDRA_FULL_ERROR":"1"},
"args": [
// "'wandb.mode=disabled",
// "env.train.id=BreakoutNoFrameskip-v4",
"env.train.id=CrafterReward-v1",
"env.train.id=Craftax-Symbolic-AutoReset-v1",
// # make it start early
"training.tokenizer.start_after_epochs=1",
"training.world_model.start_after_epochs=2",
@@ -34,8 +34,9 @@
"console": "integratedTerminal",
"justMyCode": false,
"autoReload": {"enable": true,},
"env": {"HYDRA_FULL_ERROR":"1"},
"args": [
"env.train.id=CrafterReward-v1",
"env.train.id=Craftax-Symbolic-AutoReset-v1",
]
}
]
journal.md
-196
View File
@@ -1,196 +0,0 @@
# 2023-11-03 11:36:39
Step 1 trying to get a VAE working
ah bitsandbytes
so I needed to use a `poetry add https://github.com/TimDettmers/bitsandbytes/releases/download/0.41.0/bitsandbytes-0.41.0-py3-none-any.whl` to get it to work
# 2023-11-09 07:29:27
Where am I up to?
- [ ] I want to get bigVAE working with Mistral
- [ ] run
- [ ] Then simplify it
- [ ] then understand IRIS, and inser this model
Eror: OOM w minstral. It runs out of mem when set_adapter, with an input
- 7.6/24 at first stop point we have
- vae: DecoderOnlyTransformerVAE
- self.model
- and we seem to be running inputs through with grad
- which is werid as it's jsut
- a frozen model
- 2 adaptors (float32)
- and a vae head
- `next(iter(self.model.parameters()))` is bfloat16, cuda:0. As is vae
$ self.model
PeftModel(
(base_model): LoraModel(
(model): MistralForCausalLM(
(model): MistralModel(
(embed_tokens): Embedding(32000, 4096)
$ self.vae
DecoderOnlyTransformerVAE(
(model): PeftModel(
(base_model): LoraModel(
(model): MistralForCausalLM(
(model): MistralModel(
(embed_tokens): Embedding(32000, 4096)
(vae): VAEComponent(
(f): Linear(in_features=4096, out_features=1, bias=True)
(w_e): Linear(in_features=4096, out_features=768, bias=True)
(w_d): Linear(in_features=768, out_features=4096, bias=True)
)
)
Hm JDP said he uses 8xH100 so 64GB*8.
- A p5.48xlarge is 80$
- a o2.16x is 192GB and $14/h
we are using peft
```py
self.model.print_trainable_parameters()
# trainable params: 83,886,080 || all params: 7,577,276,416 || trainable%: 1.107074302091819
self.model
```
DecoderOnlyTransformerVAE(
(model): PeftModel(
(base_model): LoraModel(
(model): MistralForCausalLM(
(model): MistralModel(
(embed_tokens): Embedding(32000, 4096)
(layers): ModuleList(
(0-31): 32 x MistralDecoderLayer(
(self_attn): MistralAttention( )
(mlp): MistralMLP()
(input_layernorm): MistralRMSNorm()
(post_attention_layernorm): MistralRMSNorm()
)
)
(norm): MistralRMSNorm()
)
(lm_head): Linear(in_features=4096, out_features=32000, bias=False)
)
)
)
(vae): VAEComponent(
(f): Linear(in_features=4096, out_features=1, bias=True)
(w_e): Linear(in_features=4096, out_features=768, bias=True)
(w_d): Linear(in_features=768, out_features=4096, bias=True)
)
)
How much gpu ram should it take it train a 7B?
- 3B -> 14
- 12B - 56GB
- so 7B should be ~30 :(. or 20 with cpu offloading
- batch size of 4... with 1 it takes 20GB then crashes
# 2023-11-10 11:41:21
tldr:
- I can't use Mistral without a bigger gpu: 30GB+, or maybe I can use the deepspeed gpu offloading (batch=1)
- I I can just use GPT2 like adavae. Or stablelm
Ideally I can use a small one for prototyping, and change to a large one if it works.
Models:
- mistral: the BigVAE code is setup for it
- gpt2: the AdaVAE code is setup for it... but it's also way messier. Roll your own adapter etc
TODO use no_grad! way better
now
- look at training code
- what is JP code actually doing? start with generate topic
- clean it up, there is so much repeeated code
# 2023-11-11 10:09:50
## Now how does VAE generate work inside?
The VAE generate function takes in context, input embeddings, target embeddings, and other parameters in that order. The function performs the following steps:
- Encode: The input embeddings are encoded using an adapter on `self.model` and `vaecomponent.encoder`. This step involves using linear and pooling operations, followed by a softmax function to sample and obtain `z`.
- Decode/Embed: `z` is decoded to obtain `z_embed`, and the decoder is provided with target ID embeddings. This step is done using a linear operation and the model embeddings.
- Model: The model is run on the embeddings to obtain logits. The embeddings can be considered as latent states, and the latent state is expressed in the language of embeddings.
# 2023-11-12 08:06:25
Now look at IRIS and TWM world models.
- https://github.com/eloialonso/iris/blob/main/src/models/world_model.py
- Model(x, kv_cache). Where x is B, T, C. Batch, Time, Channels?
- if just takes in x and outputs x. The output are logits, from a linear layer.
- https://github.com/jrobine/twm/blob/main/twm/world_model.py
# 2023-11-12 08:46:46
Adding BigVAe as transformer layer
I can't tokenize, then pass in input id's. As I need to to be backpropable. So I need to by pass the embedding layers...
- perhaps I can encode actions by things I have previously embedded?
# 2023-11-12 10:10:01
What the diff between BigVAERouter and DecoderOnlyTransformerVAE
- DecoderOnlyTransformerVAE(prefix_ids, input_ids) -> outputs, mean
- BigVAERouter(prefix_ids, input_ids, target_ids,),
OK I want to rename here:
- prefix_ids
- embed_ids -> input_ds
- target_ids
- decoder_prefix_ids -> prefix_ids
Now how to reconcille it with the world model. what does the world model do?
- components
- transforer x->x: the part we are replacing
- embedder: a custom embedder
- heads for each output: obs, reward, end_of_episode
- when it goes foward
- tokens -> (x, obs, rewards, ends)
- where is is the output of the transformer/vae
- x = transformer(sequences)
- where the sequences are tokens embedded with the embedder.
- TODO: I will want to use the model embedder if possible?
## Embedder deep dive
First the code
So the embedder, takes in `tokens` then breaks it up using a slices, into a seperate obs and action embedder.
We add them, plus positions, and add them into one of embed_dim
self.pos_emb = nn.Embedding(config.max_tokens, config.embed_dim)
self.embedder = Embedder(
max_blocks=config.max_blocks,
block_masks=[act_tokens_pattern, obs_tokens_pattern],
embedding_tables=nn.ModuleList([nn.Embedding(act_vocab_size, config.embed_dim), nn.Embedding(obs_vocab_size, config.embed_dim)])
)
sequences = self.embedder(tokens, num_steps, prev_steps) + self.pos_emb(prev_steps + torch.arange(num_steps, device=tokens.device))
Now how does the paper describe it? https://openreview.net/pdf?id=vhFu1Acb0xb
# IDEAS
Could I just use a single transformer? Dreamer seem use encode and decode to a latent state, but IRIS doesn't?
journal.md
Symlink
+1
View File
@@ -0,0 +1 @@
/home/wassname/Documents/syncthing/Markdown_notes/2024/04/27/bigvae_transformer_wm_research_log.md
justfile
+17
View File
@@ -1,11 +1,28 @@
set shell := ["zsh", "-cu"]
# Export all just variables as environment variables.
set export
WANDB_MODE := "offline"
WANDB_SILENT := "true"
HYDRA_FULL_ERROR := "1"
breakout:
. ./.venv/bin/activate
python src/main.py env.train.id=BreakoutNoFrameskip-v4
crafter:
. ./.venv/bin/activate
python src/main.py env.train.id=CrafterReward-v1
craftax:
. ./.venv/bin/activate
# https://github.dev/MichaelTMatthews/Craftax/blob/fbe4b50b985d980ea2035aa046174fb069d0cffe/craftax/craftax_env.py#L19
# python src/main.py env.train.id=Craftax-Pixels-AutoReset-v1
python -m pdb src/main.py env.train.id=Craftax-Symbolic-AutoReset-v1
# minihack:
# python src/main.py env.train.id=MiniHack-River-v0
notebooks/mjc-003_scratch_craftax.ipynb
+450
View File
File diff suppressed because one or more lines are too long
poetry.lock
Generated
+1804 -897
View File
File diff suppressed because it is too large Load Diff
pyproject.toml
+8 -4
View File
@@ -7,14 +7,14 @@ readme = "README.md"
[tool.poetry.dependencies]
python = ">=3.9,<3.13"
torch = {version = "^2.1.0+cu118", source = "pytorch"}
torch = {version = "^2.1.0+cu121", source = "pytorch"}
matplotlib = "^3.8.0"
loguru = "^0.7.2"
einops = "^0.3.1"
torchinfo = "^1.8.0"
accelerate = "^0.24.1"
peft = "^0.5.0"
bitsandbytes = {url = "https://github.com/TimDettmers/bitsandbytes/releases/download/0.41.0/bitsandbytes-0.41.0-py3-none-any.whl"}
bitsandbytes = "^0.41.0"
transformers = "4.34.0"
tqdm = "^4.66.1"
wandb = "^0.12.6"
@@ -26,14 +26,18 @@ opencv-python = "^4.8.1.78"
hydra-core = "^1.3.2"
torchvision = "^0.16.0"
numpy = ">=1.18.0"
gym = {version = "0.22.0", extras = ["accept-rom-license", "atari"]}
scipy = "^1.11.3"
crafter = "^1.8.2"
minihack = "^0.1.5"
craftax = {path = "/media/wassname/SGIronWolf/projects5/2024/Craftax"}
gymnasium = "^0.29.1"
jax = "^0.4.28"
jaxlib = "^0.4.28"
pandas = "^2.2.2"
[[tool.poetry.source]]
name = "pytorch"
url = "https://download.pytorch.org/whl/cu118"
url = "https://download.pytorch.org/whl/cu121"
priority = "explicit"
[tool.poetry.group.dev.dependencies]
research_journal.md
-407
View File
@@ -1,407 +0,0 @@
# 2023-11-12 13:17:35
Try IRIs but with pretrained transformer with LoRA adapter
- [x] first can I run it yes with a 1/2 batch size
- [ ] then can I add 3B with adapter...
```sh
poetry install
. ./.venv/bin/activate
python src/main.py env.train.id=BreakoutNoFrameskip-v4 common.device=cuda:0 wandb.mode=offline
# or for quick debug
WANDB_MODE=disabled python -m pdb src/main.py env.train.id=BreakoutNoFrameskip-v4
```
```sh
# TODO use this code to load a transformer, and other code from my bigvae repo https://github.com/wassname/bigvae_wm
def load_model(config, device='cuda'):
tokenizer = AutoTokenizer.from_pretrained(config.model_name, trust_remote_code=True)
tokenizer.padding_side = "left"
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True,
)
base_model = AutoModelForCausalLM.from_pretrained(
config.model_name,
device_map={"": device},
quantization_config=bnb_config,
torch_dtype=torch.bfloat16,
trust_remote_code=True
)
peft_config = peft.LoraConfig(
peft.TaskType.CAUSAL_LM,
inference_mode=False,
r=config.rank,
lora_alpha=8,
lora_dropout=config.dropout,
target_modules=[
"self_attn.q_proj",
"self_attn.k_proj",
"self_attn.v_proj",
"self_attn.o_proj",
"mlp.gate_proj",
"mlp.up_proj",
"mlp.down_proj",
],
)
base_model_peft = peft.get_peft_model(base_model, peft_config)
vae_model = BigVAE(
base_model_peft, device, peft_config, z_dim=config.z_dim,
)
if config.start_from:
vae_model.load_pretrained(config.start_from)
base_model_peft.requires_grad_(False)
vae_model.vae_head.requires_grad_(False)
vae_model.vae_head.w_d.requires_grad_()
router = BigVAERouter(base_model_peft, vae_model, device, peft_config)
if config.start_from:
router.load_pretrained(config.start_from, is_trainable=True)
print(router.model.print_trainable_parameters())
router.model.set_adapter("router")
```
Debugging:
batch['observations'].shape
torch.Size([16, 20, 3, 64, 64])
obs_tokens.shape
torch.Size([16, 20, 16])
https://vscode.dev/github/wassname/iris_bigvae/blob/just_llms2/src/models/world_model.py#L105
tokens
tensor([[222, 222, 222, ..., 409, 55, 2],
[222, 222, 222, ..., 409, 139, 1],
[222, 222, 222, ..., 168, 190, 3],
...,
[222, 222, 222, ..., 168, 55, 0],
[222, 222, 222, ..., 237, 190, 3],
[222, 222, 222, ..., 168, 55, 0]], device='cuda:0')
tokens.shape
torch.Size([16, 340])
where 16 is the batch size. 340 is the step size?. actions was 16,20 int
tokens.shape int
torch.Size([16, 340])
sequences.shape float32
torch.Size([16, 340, 256])
transfrmer
x.shape
torch.Size([16, 340, 256])
# 2023-11-12 16:58:37
So I got it training, but during imagination it passes in a single token with no past steps. But the slicer seems to need at least on block? And so I get none?
hmm it's because num_kept_tokens is 16 not 1. So there should be a whole block passed in ?
wait apparently it's also a problem in the normal repo.... I confuse! maybe it's my config! maybe I need >larger than block size. nope
hmm it still happens in the original repo with my debug params. maybe it's my debug params
... trying a full run without my debug params...
note trains.world_model.batch_num_samples:4 fill 20GB gpu ram for the 3b stability ai llm
ok even with a full run I get the error. I think it's a bug in the original repo. I'll try to debug it there.
Epoch 51 / 600
Experience collection (train_dataset): 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:03<00:00, 59.91it/s]
Training tokenizer: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:17<00:00, 11.53it/s]
Training world_model: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [02:11<00:00, 1.53it/s]
Training actor_critic: 0%| | 0/200 [00:00<?, ?it/s]
Error executing job with overrides: ['env.train.id=BreakoutNoFrameskip-v4', 'common.device=cuda:0', 'wandb.mode=offline']
Traceback (most recent call last):
File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/src/main.py", line 10, in main
trainer.run()
File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/src/trainer.py", line 111, in run
to_log += self.train_agent(epoch)
File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/src/trainer.py", line 146, in train_agent
metrics_actor_critic = self.train_component(self.agent.actor_critic, self.optimizer_actor_critic, sequence_length=1 + self.cfg.training.actor_critic.burn_in, sample_from_start=False, tokenizer=self.agent.tokenizer, world_model=self.agent.world_model, **cfg_actor_critic)
File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/src/trainer.py", line 161, in train_component
losses = component.compute_loss(batch, **kwargs_loss) / grad_acc_steps
File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/src/models/actor_critic.py", line 102, in compute_loss
outputs = self.imagine(batch, tokenizer, world_model, horizon=imagine_horizon)
File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/src/models/actor_critic.py", line 149, in imagine
obs, reward, done, _ = wm_env.step(action_token, should_predict_next_obs=(k < horizon - 1))
File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/src/envs/world_model_env.py", line 75, in step
reward = Categorical(logits=outputs_wm.logits_rewards).sample().float().cpu().numpy().reshape(-1) - 1 # (B,)
File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/torch/distributions/categorical.py", line 70, in __init__
super().__init__(batch_shape, validate_args=validate_args)
File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/torch/distributions/distribution.py", line 66, in __init__
valid = constraint.check(value)
File "/media/wassname/SGIronWolf/projects5/worldmodels/iris_bigvae/.venv/lib/python3.9/site-packages/torch/distributions/constraints.py", line 226, in check
result = result.reshape(
RuntimeError: cannot reshape tensor of 0 elements into shape [8, 0, -1] because the unspecified dimension size -1 can be any value and is ambiguous
Oh maybe it's because we don't keep track of KV cache, but it's actually used to track number of steps!!
# 2023-11-13 20:11:51
I go it working byt ut takes 30 seconds for one one actor critic batch, werird
Experience collection (train_dataset): 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:03<00:00, 60.45it/s]
Training tokenizer: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:17<00:00, 11.26it/s]
Training world_model: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [02:12<00:00, 1.51it/s]
Training actor_critic: 82%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▍ | 165/200 [1:01:43<13:24, 22.99s/it]
hm maybe it's just the face it has to backprop throguh the whole LLM :( damn... is there another way to train it? Daym. How many params did the original have?
well running eval on the transformer brought it down from 100sec to 60, but it's still huge.
But then why is the model training fast? It makes not sense
# 2023-11-16 12:54:48
Why is agent so slow? Lets find out
- look at diagram
- look at train_agent
- to tokenizer.compute_loss is just tokenizer
- world_model.compute_loss user tokenizer with no grad
- actor_critic? takes an hour!!
- imagine (with grad?)
- x20 = horizon
- self(obs)
- WorldModelEnv.step this has no grad!
- transformer
- tokenizer with no grad
- compute_lambda_returns with no grad
So changes:
- the world model step always had no grad!
- I just made the lstm smaller and the horizon smaller
- from 1h to 3m. Reasonable.
Experiment:
- try no grad on the model? ok it now takes 20 minutes to train... still slow
with a smaller lstm and only 10 steos ut tajes 8 mins,
![](img/2023-11-16-13-01-11.png)
Experience collection (train_dataset): 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:05<00:00, 36.11it/s]
Training tokenizer: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:55<00:00, 3.61it/s]
Training world_model: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [01:12<00:00, 2.76it/s]
Training tokenizer: 55sec
Training world_model 72 sec
train actor_critic 3min. It looks like it scales with lstm size!
new changes 10mins
lets try no lstm?
Right now it will take 41 hours for on epoch lol
ram during stages:
- actor critic 20G/24
how big does my actor critic need to be?
- IRIS: large 512 lstm on 64,64,3 obs
- We ran our experiments with 8 Nvidia A100 40GB GPUs. With two Atari environments running on the same GPU, training takes around 7 days, resulting in an average of 3.5 days per environment.
- twm: mlp 512
How long to train?
`600*10//6/24` = 41 days
- 600 epochs * 10 minutes / 6 to get hours, 24 to get days
# 2023-11-17 07:59:44
so I've got it working with these times. But maybe it's too small
Epoch 148 / 600
Experience collection (train_dataset): 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:10<00:00, 19.25it/s]
Training tokenizer: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:59<00:00, 3.34it/s]
Training world_model: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [01:15<00:00, 2.64it/s]
Training actor_critic: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 20/20 [03:20<00:00, 10.05s/it]
- what about resume? oh we seem to have that although the code doesn't make sense https://hydra.cc/docs/tutorials/basic/running_your_app/working_directory/ https://hydra.cc/docs/1.2/upgrades/1.1_to_1.2/changes_to_job_working_dir/ see eval.py
- [ ] but it's still too damn slow. what about bfloat16? using auto case?
- why does it take so long? it would be nice to have a reproduction notebook
- also the model might be to small now....
how to play
```sh
cd outputs/2023-11-17/07-59-44
python scripts/play.sh
```
## Envs
tl:dr just use pong or breakout or crafter (1m steps)
for steps see [crafter paper](https://arxiv.org/pdf/2109.06780.pdf)
Nethack learning env. What's the obs size? 21x79 of glyphs (5991 possibilities) and 21 dim of stats
- they use an lstm of 128. 5 layer conv
- requires 1B steps
-
atari:
- reqs 200M stpes
progcen:
- 200M steps
minihack:
- 2M steps for room 5xt
- but needs editing to be atari compatible. e.g. 336 × 1264 × 3 pixels
- pixel_crop 64,64,3 or 9x9 crop works!
- lstm 256
- The training on MiniHack’s Room-5x5 task for two million timesteps using our IMPALA baseline takes approximately 4:30 minutes (r
crafter
- reqs 1M steps
- "All agents trained for 1M environment steps in under 24 hours on a single GPU and we repeated the training for 10 random seeds per method. The training reward curves are included in Appendix "
from https://arxiv.org/pdf/2111.09794.pdf
There are several PCG state-varying gridworld environments (
- [MiniGrid](https://minigrid.farama.org/environments/minigrid/), ~~- BabyAI~~
- Crafter,
- 2019 [Rogue-gym,](https://github.com/kngwyu/rogue-gym)
- 2020 MarsExplorer, maxe exploration. 1M steps
- NLE,
- MiniHack;
- [gym\_nethack](http://campbelljc.com/research/gym_nethack/)
- 2018 [rogueinabox](https://github.com/rogueinabox/rogueinabox)
- [rogue-gym](https://github.com/kngwyu/rogue-gym)
- [MiniGrid](https://github.com/maximecb/gym-minigrid)
- 2019 [CoinRun](https://github.com/openai/coinrun) no traction or maintanance
- [MineRL](http://minerl.io/docs)
- [Project Malmo](https://www.microsoft.com/en-us/research/project/project-malmo/) miencraft
- [OpenAI Procgen Benchmark](https://openai.com/blog/procgen-benchmark/) 200M steps
- 2020 [Obstacle Tower](https://github.com/Unity-Technologies/obstacle-tower-env) - 3d slow
non-PCG observation-varying continuous control environments
- (RoboSuite, DMC-Remastered, DMC-GB, DCS, KitchenShift, NaturalEnvs MuJoCo; Fan
et al., 2021; Grigsby & Qi, 2020; Hansen & Wang, 2021; Stone et al., 2021; Xing et al., 2021a;
Zhang et al., 2018a), and multi-task continuous control benchmarks which could be adapted
to ZSG (CausalWorld, RLBench, Meta-world; Ahmed et al., 2020; James et al., 2019a; Yu
et al., 2019).
## investigating model slowness
So it's all just that using the transformer to imagine takes almost 0.1s. but it's run so many more times than during training. All my ideas to speed it up don't work.
- [x] eval. no grad
- [x] remove the call for adapter, causal mask each time?
- [ ] lower rank?
Ok so it's all just the
- rollout, controller by max block size. 10x
- the fact that actor_critic can use a larger batch, therefore 4-8x more samples
- for each one it imagines 2
# 2023-11-18 06:17:55
It trained overnight, now I would like to view a replay
Hmm "delta-IRIS" ∆-IRIS
https://openreview.net/forum?id=o8IDoZggqO
∆-IRIS encodes
new frames by attending to the ongoing trajectory, effec-
tively describing deltas between timesteps.
This new ap-
proach drastically reduces the number of tokens to encode
frames, since they are not encoded independently as in IRIS.
In the Crafter benchmark (Hafner, 2022), ∆-IRIS unlocks
16 out of 22 objectives at the 10M frames mark
# 2023-11-18 16:16:16
Why is it not learning? It's because the dynamics model is total BS!!!
- [ ] Well lets try training it for longer then. It's cheap to train so..
- [ ] also maybe train tokenizer and model together? I have a lot of frozen layers, including the embeddings... so might be better
- [ ] oh no we do have an unforzen embedder before the transformer or more layers
- maybe I need a higher rank lora? after all I'm changing a lot from text tokens
- maybe no tokens, bypass to embedder?
# 2023-11-19 06:45:34
So I tried just trainign the world model for 200 epochs. And with a post_embedding layer. It helped the flickering. But not enougth to actually go for the obvious local minima of the next state equals the last
idea
- bypass embedding?, but wait dreamerv3 needed quant z...
- yes I am bypassing it by passing in the input_embeds... but maybe I shouldn't
- [x] use same embedding everywhere. e.g. model embedding in encoder decoder?
- Our embedings is (embed_tokens): Embedding(32000, 2048). So we would need to encode to 32000!
ok we need to freeze it, and change dtype
OK it seems slightly better yay! Lets train it overnight and see
next idea is to the delta-IRIS thing where the tokens only have to encode the diff(obs)
# 2023-11-19 16:50:45
Seems to be working! Now let's plan delta-IRIS
So IRIS has
- Encoder $E(x_0, a_0) = t_0$
```py
obs_tokens = self.tokenizer.encode(observations, should_preprocess=True).tokens # (B, C, H, W) -> (B, K)
```
- Embed $Emb(t_0) = z_0$
```py
embedded_tokens = self.tokenizer.embedding(self.obs_tokens) # (B, K, E)
z = rearrange(embedded_tokens, 'b (h w) e -> b e h w', h=int(np.sqrt(self.num_observations_tokens)))
```
- Dynamics $D(z_0, a_0) = z_1$
```py
outputs_wm = self.world_model(tokenRedmond AI, past_keys_values=self.keys_values_wm)
```
- Decoder $D(z_0, a_0) = x_1$
```py
rec = self.tokenizer.decode(z, should_postprocess=True) # (B, C, H, W)
```
but we have tokens vs z
Questions:
- wait why are we just passing in "action_token" to the transformer and not obs? that must have obs in it right... right??? confirm
- in iris-delta how did they pass everything in? I guess obs_prev was tokenized too? I think the slices are annoying so maybe I should just pass things seperatly
# 2023-11-24 10:56:40
If I unfreeze the whole transformer, it seem to learn the most obvious dynamics (the next latent space is the same as the last).
To summarize
- with Qlora it didn't learn that
- with unfrozen head it didn't
- when training transformer and obs embedding together it did not (frozen llm embeddings)
no it didn't work with tokenizer sep hmm
Oh it did with whole transfrmer and tokenizer at same time https://wandb.ai/wassname/iris/runs/w7lvs4gi?workspace=user-wassname
wandb: world_model/eval/loss_obs ▇█▃▄▄▄▄▃▃▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁
wandb: world_model/eval/loss_rewards ▁█▃█▇▅▇▁▅▆▂▄▃▇▆▃▄▇▃▄▂▄▃▂▃▄▃
wandb: world_model/eval/total_loss ▂█▃██▅▇▁▅▆▁▃▂█▇▃▃▆▃▄▂▃▃▂▂▄▂
src/envs/wrappers.py
+49 -1
View File
@@ -6,18 +6,66 @@ from typing import Tuple
import gym
import numpy as np
import jax
from PIL import Image
import crafter
from craftax.craftax_env import make_craftax_env_from_name
from craftax.craftax.play_craftax import CraftaxRenderer
def make_env(id, size=64, max_episode_steps=None, noop_max=30, frame_skip=4, done_on_life_loss=False, clip_reward=False):
if id.startswith('Craftax'):
return make_craftax(id)
if id.startswith('Crafter'):
return make_crafter(id, size=size, max_episode_steps=max_episode_steps, done_on_life_loss=done_on_life_loss)
if id.startswith('MiniHack'):
return make_minihack(size=size, max_episode_steps=max_episode_steps, done_on_life_loss=done_on_life_loss)
else:
return make_atari(id, size, max_episode_steps, noop_max, frame_skip, done_on_life_loss, clip_reward)
class Gymax2GymWrapper(gym.ObservationWrapper):
def __init__(self, env) -> None:
gym.Wrapper.__init__(self, env)
self.env_params = env.default_params
rng = jax.random.PRNGKey(0)
rng, _rng = jax.random.split(rng)
self.rngs = jax.random.split(_rng, 3)
self.do_render = True
if self.do_render:
self.renderer = CraftaxRenderer(self.env, self.env_params, pixel_render_size=1)
def step(self, action):
obs, state, reward, done, info = self.env.step(self.rngs[2], self.env_state, action, self.env_params)
self.env_state = state
if self.do_render:
self.renderer.update()
return obs, reward, done, info
@property
def action_space(self):
return self.env.action_space(self.env_params)
@property
def observation_space(self):
return self.env.observation_space(self.env_params)
def reset(self):
obs, state = self.env.reset(self.rngs[0], self.env_params)
self.env_state = state
if self.do_render:
self.renderer.update()
return obs
def render(self):
self.renderer.render(self.env_state)
def make_craftax(id, size=64, max_episode_steps=None, done_on_life_loss=False):
env = make_craftax_env_from_name(id)
env = Gymax2GymWrapper(env) # # (130, 110, 3) or (9, 11, 83)
# if 'pixel' in id:
# env = ResizeObsWrapper(env, (size, size))
return env
def make_atari(id, size=64, max_episode_steps=None, noop_max=30, frame_skip=4, done_on_life_loss=False, clip_reward=False):
env = gym.make(id)
@@ -38,7 +86,7 @@ def make_atari(id, size=64, max_episode_steps=None, noop_max=30, frame_skip=4, d
def make_crafter(id, size=64, max_episode_steps=None, done_on_life_loss=False):
# https://github.com/danijar/dreamerv2/blob/07d906e9c4322c6fc2cd6ed23e247ccd6b7c8c41/dreamerv2/common/envs.py#L242
# https://github.com/footoredo/torchbeast/blob/12939569cc46b6a8616e4c25b138d97248cc8581/torchbeast/atari_wrappers.py#L301
env = gym.make(id)
env = gym.make(id)
env = ResizeObsWrapper(env, (size, size))
return env
src/main.py
+1
View File
@@ -4,6 +4,7 @@ from omegaconf import DictConfig
from trainer import Trainer
from loguru import logger
import sys
logger.add(sys.stderr, format="{time} {level} {message}", filter="my_module", level="INFO")
@hydra.main(config_path="../config", config_name="trainer")
src/models/bigvae.py
+7 -6
View File
@@ -21,7 +21,7 @@ import sys
import zipfile
import accelerate
from datasets import load_dataset
# from datasets import load_dataset
import peft
import safetensors.torch as safetorch
import torch
@@ -347,11 +347,6 @@ def batched(iterable, n):
@dataclass
class BigVAEConfig:
model_name: str = "stabilityai/stablelm-3b-4e1t"
dropout: float = 0
rank: int = 32
z_dim: int = 768
start_from: str = None
tokens_per_block: int
max_blocks: int
@@ -364,6 +359,12 @@ class BigVAEConfig:
embed_pdrop: float
resid_pdrop: float
attn_pdrop: float
model_name: str = "stabilityai/stablelm-3b-4e1t"
dropout: float = 0
rank: int = 32
z_dim: int = 768
start_from: str = None
@property
def max_tokens(self):
src/models/transformer.py
+1 -1
View File
@@ -12,7 +12,7 @@ import torch.nn as nn
from torch.nn import functional as F
from loguru import logger
# from .kv_caching import KeysValues, KVCache
from .kv_caching import KeysValues, KVCache
@dataclass
src/models/world_model.py
+1
View File
@@ -12,6 +12,7 @@ from src.models.slicer import Embedder, Head
from src.models.tokenizer import Tokenizer
from src.models.transformer import Transformer, TransformerConfig
from src.utils import init_weights, LossWithIntermediateLosses
from src.models.bigvae import BigVAEConfig
@dataclass