Skip to content
Open

4.0 #402

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
298 commits
Select commit Hold shift + click to select a range
d319dbc
merge
jsuarez5341 Oct 8, 2025
86dcf77
Impulse wars config
jsuarez5341 Oct 9, 2025
f94dc05
nmmo3 config
jsuarez5341 Oct 9, 2025
2d21214
cogames
jsuarez5341 Oct 9, 2025
bcb6a9f
nmmo3 net params
jsuarez5341 Oct 9, 2025
af07504
fuck this
capnspacehook Aug 20, 2025
d6b2eab
Fix impulse wars
jsuarez5341 Oct 9, 2025
c42db1c
IW sweep
jsuarez5341 Oct 9, 2025
c32430f
initial box plot
jsuarez5341 Oct 9, 2025
d85672a
merge
jsuarez5341 Oct 9, 2025
492ec08
To pandas
jsuarez5341 Oct 9, 2025
80c5712
progress
jsuarez5341 Oct 10, 2025
3e6dd21
constellation cache
jsuarez5341 Oct 11, 2025
6036f81
Prototype in color
jsuarez5341 Oct 11, 2025
2b0d829
UI
jsuarez5341 Oct 11, 2025
dad30c9
prototype
jsuarez5341 Oct 11, 2025
6f67327
pretty
jsuarez5341 Oct 12, 2025
6503dbd
Progress!
jsuarez5341 Oct 13, 2025
854c3ec
Progress!
jsuarez5341 Oct 13, 2025
9077b0e
shaders!
jsuarez5341 Oct 14, 2025
80d8127
UI
jsuarez5341 Oct 15, 2025
ce29f53
UI
jsuarez5341 Oct 15, 2025
bfc2862
shaders
jsuarez5341 Oct 15, 2025
aa4f90f
constellations!
jsuarez5341 Oct 15, 2025
c746d1a
minor refactor
jsuarez5341 Oct 15, 2025
48c5329
Tooltip prototype
jsuarez5341 Oct 16, 2025
5c7e270
Initial refactor
jsuarez5341 Oct 17, 2025
33bd507
tooltip
jsuarez5341 Oct 17, 2025
21d66c7
Latest
jsuarez5341 Oct 18, 2025
740eee3
latest
jsuarez5341 Oct 18, 2025
65cdba0
UNSTABLE TESTING. DO NOT USE
jsuarez5341 Oct 18, 2025
f168bd6
BROKEN- DO NOT USE
jsuarez5341 Oct 18, 2025
9d23693
Initial cuda bind
jsuarez5341 Oct 19, 2025
62d361f
Build flags
jsuarez5341 Oct 20, 2025
3452174
comprehensive perf test
jsuarez5341 Oct 23, 2025
883a3d8
Iniital puffer cpp
jsuarez5341 Oct 23, 2025
81de3c5
4.6M sps test
jsuarez5341 Oct 24, 2025
4d3cdbd
Un-cppify a bit
jsuarez5341 Oct 24, 2025
3699447
Cosine anneal lr
jsuarez5341 Oct 24, 2025
8f64e12
bf16. It is slower. Will fix after
jsuarez5341 Oct 24, 2025
bcef486
Add logs
jsuarez5341 Oct 27, 2025
c68c4c6
Progress on port -- need a full matching test from python to Cpp for …
jsuarez5341 Oct 27, 2025
d1395f1
Reproducible net init
jsuarez5341 Oct 28, 2025
c0b6cf5
almost there!
jsuarez5341 Oct 30, 2025
eeba3ed
Assert check passes
jsuarez5341 Oct 30, 2025
d38cec6
fixed the bug! Was the lstm shared init for cpp
jsuarez5341 Oct 30, 2025
e94ea35
CPU check pass
jsuarez5341 Oct 31, 2025
9c4fef7
Enable gpu. Check fails, trains, 4msps)
jsuarez5341 Oct 31, 2025
6ca97e0
bf16. Slower for now because missing kernel. May need to do 32b accum
jsuarez5341 Oct 31, 2025
5b0638c
Mingru
jsuarez5341 Nov 4, 2025
0b50ca1
Multilayer
jsuarez5341 Nov 4, 2025
d7f8880
progress
jsuarez5341 Nov 4, 2025
a8d1423
Add torch muon
jsuarez5341 Nov 4, 2025
3074858
Merge pull request #411 from PufferAI/muon
jsuarez5341 Nov 4, 2025
bc74c3f
Stable
jsuarez5341 Nov 4, 2025
676bdf2
Update sweep defaults
jsuarez5341 Nov 4, 2025
0575058
Mamba
jsuarez5341 Nov 5, 2025
faede5f
Merge branch '4.0' into mamba
jsuarez5341 Nov 5, 2025
bde228c
Merge pull request #412 from PufferAI/mamba
jsuarez5341 Nov 5, 2025
9df8eda
Testing new archs
jsuarez5341 Nov 5, 2025
35530d5
Minor
jsuarez5341 Nov 6, 2025
6efab69
Ready for sweeps
jsuarez5341 Nov 6, 2025
86bc81c
Initial mingru
jsuarez5341 Nov 6, 2025
304990c
Update configs for sweeping
jsuarez5341 Nov 7, 2025
a9926ce
merge
jsuarez5341 Nov 7, 2025
01071cb
Adam for now so we can run
jsuarez5341 Nov 7, 2025
84636a7
Initial kern
jsuarez5341 Nov 8, 2025
47e4d24
merge
jsuarez5341 Nov 8, 2025
de94e75
merge
jsuarez5341 Nov 8, 2025
0f67928
More kernels
jsuarez5341 Nov 8, 2025
effe468
test kerns
jsuarez5341 Nov 8, 2025
47a0471
decent kernels
jsuarez5341 Nov 8, 2025
5d5655a
Progress
jsuarez5341 Nov 9, 2025
77011ed
Kernels pass
jsuarez5341 Nov 9, 2025
018da4d
Fix cast
jsuarez5341 Nov 9, 2025
f1ef5e2
comment prints
jsuarez5341 Nov 9, 2025
0b7f7b4
Perfect logcoeff kernel
jsuarez5341 Nov 11, 2025
456b56f
Numerically stable fused scan
jsuarez5341 Nov 11, 2025
ba446d5
Stable train
jsuarez5341 Nov 12, 2025
b22b41d
Fix bias in entropy grad
jsuarez5341 Nov 12, 2025
6f2bd90
6.6m w/ cpu envs
jsuarez5341 Nov 12, 2025
0b63fbb
skip connect and rmsnorm makes mingru stable
jsuarez5341 Nov 13, 2025
e2cbb3e
Fix breakout; solid mingru in python, runnable cpp
jsuarez5341 Nov 13, 2025
b86792c
multilayer
jsuarez5341 Nov 13, 2025
91a0e9f
latest
jsuarez5341 Nov 14, 2025
7eb5799
nmmo3
jsuarez5341 Nov 14, 2025
fa76d1b
Initial muon (needs binds)
jsuarez5341 Nov 14, 2025
0838c91
Merge branch '4.0' of https://github.com/pufferai/pufferlib into 4.0
jsuarez5341 Nov 14, 2025
6fbc382
RMSNorm
jsuarez5341 Nov 15, 2025
52920d6
Initial dll vec
jsuarez5341 Nov 15, 2025
ef0ab9a
DLL-based training initial
jsuarez5341 Nov 15, 2025
6dbb80f
Merge branch '4.0' into 4.0-merge
jsuarez5341 Nov 17, 2025
ba605bc
Merge pull request #419 from PufferAI/4.0-merge
jsuarez5341 Nov 17, 2025
60c6e62
merge
jsuarez5341 Nov 17, 2025
da8334a
Merge branch '4.0' of https://github.com/pufferai/pufferlib into 4.0
jsuarez5341 Nov 17, 2025
1f7da1a
ready g2048 for testing
jsuarez5341 Nov 17, 2025
35dfaeb
fix optim and config
jsuarez5341 Nov 17, 2025
a6cd57a
Tweak torch muon
jsuarez5341 Nov 17, 2025
a55b5e9
Fix sweeps
jsuarez5341 Nov 17, 2025
4219ee9
Torch muon matched to heavyball muon numerics
jsuarez5341 Nov 18, 2025
4a143b1
Migrate optimizer to custom muon
jsuarez5341 Nov 18, 2025
b24f130
Fix big layers
jsuarez5341 Nov 18, 2025
804fa1c
Initial heavyball muon to cpp
jsuarez5341 Nov 18, 2025
e529fb6
Working initial vec
jsuarez5341 Nov 20, 2025
6e973cc
Initial pinned mem
jsuarez5341 Nov 20, 2025
42a5b3d
Progress on kerns
jsuarez5341 Nov 22, 2025
7904c0b
latest
jsuarez5341 Nov 25, 2025
abd7965
latest
jsuarez5341 Nov 25, 2025
87f0529
Merge branch '4.0' into merge
jsuarez5341 Nov 26, 2025
685933e
Merge pull request #427 from PufferAI/merge
jsuarez5341 Nov 26, 2025
3b32ffc
temp fixes
jsuarez5341 Nov 26, 2025
51ad295
Initial cuda buffering
jsuarez5341 Nov 26, 2025
ed81f3e
Initial buffered vec. Compiles, hangs with >1 buffer
jsuarez5341 Nov 26, 2025
f0fdb6a
temp
jsuarez5341 Nov 28, 2025
e79ba75
env c files
jsuarez5341 Nov 28, 2025
e5d66a2
20k map gen for tower climb, fast load, made pufferl work
kywch Nov 28, 2025
5eb660b
Initial buffered vec runnable
jsuarez5341 Nov 28, 2025
99daf55
Merge pull request #431 from kywch/tower-20k
jsuarez5341 Nov 28, 2025
49ee5db
Blocked perf test
jsuarez5341 Nov 29, 2025
825bc9b
works 1 buffer, not 2
jsuarez5341 Nov 29, 2025
1f8d8ac
Fix bug. Am idiot
jsuarez5341 Dec 1, 2025
17f4920
latest vec
jsuarez5341 Dec 3, 2025
1437acb
initial vec
jsuarez5341 Dec 6, 2025
d2fc984
config etc
jsuarez5341 Dec 6, 2025
430b8ed
Mutex test
jsuarez5341 Dec 18, 2025
707f0f0
Progress
jsuarez5341 Dec 19, 2025
84905f8
Working initial cpu!
jsuarez5341 Dec 19, 2025
e954a9b
Now it actually works on cpu
jsuarez5341 Dec 19, 2025
9b4d8c1
Initial gpu c vec
jsuarez5341 Dec 19, 2025
0e14de0
6m sps breakout training. A bit unstable, but we got a 22s solve
jsuarez5341 Dec 19, 2025
a932085
Still unstable but good start for profiling
jsuarez5341 Dec 19, 2025
96c3b4b
Muon optims
jsuarez5341 Dec 20, 2025
0254d42
temp
jsuarez5341 Dec 20, 2025
39f1cd8
Temp
jsuarez5341 Dec 26, 2025
f15764d
temp cudagraphs
jsuarez5341 Dec 26, 2025
ec1dc6b
single buffer graph:
jsuarez5341 Dec 26, 2025
de2fe54
Working kernel. Major bug fix. Launch on same stream as torch
jsuarez5341 Dec 27, 2025
3b53176
9m with cuda trace backward
jsuarez5341 Dec 27, 2025
be99a4f
fixed fused loss for now
jsuarez5341 Dec 27, 2025
07948e9
Initial refactor
jsuarez5341 Dec 29, 2025
7f09246
Kern/graph options
jsuarez5341 Dec 29, 2025
a829b2d
pufferlib/pufferl.py
jsuarez5341 Dec 30, 2025
1a7e3c9
Initial numerical match to unoptimized cpp
jsuarez5341 Dec 30, 2025
367c9b2
stable training with cudagraphs
jsuarez5341 Dec 31, 2025
b5accce
Major env vec bug fix
jsuarez5341 Dec 31, 2025
a8e3c4a
6m stable
jsuarez5341 Dec 31, 2025
393afe0
graph rollout copy
jsuarez5341 Dec 31, 2025
0065ce3
20s with norm, 16s (unstable) without)
jsuarez5341 Dec 31, 2025
9b8f4d0
New gate!
jsuarez5341 Jan 3, 2026
f3a2c61
nvtx
jsuarez5341 Jan 6, 2026
36c790c
Initial profile design
jsuarez5341 Jan 6, 2026
4c42788
Profile full w/ cpp
jsuarez5341 Jan 6, 2026
7909a81
Initial kernel profile. Needs some renames + consistency checks. Also…
jsuarez5341 Jan 6, 2026
bc23919
defensive commit
jsuarez5341 Jan 7, 2026
540f572
Stable profiling
jsuarez5341 Jan 7, 2026
cda832a
minor profiling
jsuarez5341 Jan 7, 2026
c6b7df7
Experimental - fuse cat + pad into scan
jsuarez5341 Jan 7, 2026
3e34d43
Temp - not sure if I like this. We need to fuse a bunch of kernel ops…
jsuarez5341 Jan 8, 2026
169d521
Fuse mingru gate infrence
jsuarez5341 Jan 10, 2026
953c700
Fuse log coff and value kernel with scan
jsuarez5341 Jan 10, 2026
f30e493
Messy fuse chunk and gate to scan
jsuarez5341 Jan 10, 2026
1c27434
Fix profiler
jsuarez5341 Jan 10, 2026
932d444
9m
jsuarez5341 Jan 12, 2026
d85aca1
nprof
jsuarez5341 Jan 13, 2026
3845580
9.9m - optimized eval to 3-4 kernels
jsuarez5341 Jan 13, 2026
3262e1c
Initial 10m stable
jsuarez5341 Jan 13, 2026
a358c79
Initial profiling
jsuarez5341 Jan 15, 2026
cb58e76
Initial refactor
jsuarez5341 Jan 17, 2026
ce05b56
Refactor for deduped profiling
jsuarez5341 Jan 17, 2026
93f68a7
new bind file
jsuarez5341 Jan 17, 2026
9e71a07
more refactor
jsuarez5341 Jan 17, 2026
8f54000
Major refactor
jsuarez5341 Jan 22, 2026
e2178cf
Bind multiple envs, begin adding multidiscrete support
jsuarez5341 Jan 22, 2026
fe1bea0
initial
jonahsamost Jan 23, 2026
9eb0b97
minor
jonahsamost Jan 23, 2026
e3bbe5d
remove comments
jonahsamost Jan 23, 2026
d59bdb3
simplified 2048
kywch Jan 24, 2026
82dabc9
tweak rew
kywch Jan 24, 2026
6a56c11
multidiscrete working
jsuarez5341 Jan 24, 2026
6e85b66
simplify and edits for 132k
kywch Jan 24, 2026
ef81be6
match c network to torch net
kywch Jan 24, 2026
96c65e7
Support different num agents per env with custom vec init
jsuarez5341 Jan 24, 2026
b97d065
ppo
jonahsamost Jan 24, 2026
7cd6a23
simplified scaffolding
kywch Jan 24, 2026
705acc9
Merge pull request #472 from jonahsamost/jonah_ppo_1_24
jsuarez5341 Jan 24, 2026
e431e55
Testing ppo kernel
jsuarez5341 Jan 24, 2026
1359de8
g2048 changes only
kywch Jan 24, 2026
f0bfbd1
Merge pull request #474 from kywch/s2048-puf4
jsuarez5341 Jan 24, 2026
6f4a7f8
2048 in 4.0
jsuarez5341 Jan 24, 2026
b4ac775
Initial drive bind
jsuarez5341 Jan 25, 2026
c0b350d
Initial bindings for ~all envs
jsuarez5341 Jan 25, 2026
f2b81d9
initial
jonahsamost Jan 25, 2026
78aea9f
add TODO
jonahsamost Jan 25, 2026
dcfebd7
add strides to ppo
jonahsamost Jan 26, 2026
d77eb66
Merge pull request #475 from jonahsamost/jonah_ppo_1_25
jsuarez5341 Jan 27, 2026
f1d7917
Vec win from Peru - 11.8M sps
jsuarez5341 Jan 27, 2026
78b28bc
11m with parallel rollouts
jsuarez5341 Jan 29, 2026
f4464cc
Move graph set outside fn
jsuarez5341 Jan 29, 2026
ca54b67
Small refector
jsuarez5341 Jan 29, 2026
c72cefd
Fuse rollout graphs
jsuarez5341 Jan 30, 2026
5c4e2a6
static
jsuarez5341 Jan 30, 2026
6413f45
Initial static bind
jsuarez5341 Jan 30, 2026
142cc40
trains but not well
jonahsamost Jan 30, 2026
f68f7f7
working
jonahsamost Jan 30, 2026
0763e63
Contiguous weight update
jsuarez5341 Jan 30, 2026
e7ae2d7
Drive kernel
jsuarez5341 Jan 30, 2026
95c05d2
working both
jonahsamost Jan 30, 2026
0f028d8
Update binds
jsuarez5341 Jan 31, 2026
f027552
Merge pull request #480 from PufferAI/static
jsuarez5341 Jan 31, 2026
9e1f1b4
Jonah's bf16
jsuarez5341 Jan 31, 2026
6ad185b
Major bug fix on cuda streams, stable training
jsuarez5341 Jan 31, 2026
162b181
Stable for many envs. kernels + bf16 iffy
jsuarez5341 Jan 31, 2026
8474f2d
Cleaned up, 14m breakout
jsuarez5341 Jan 31, 2026
3df1d67
Runnable sweeps
jsuarez5341 Feb 3, 2026
beb1419
Sweeps integration + mem leak fixes
jsuarez5341 Feb 3, 2026
0f530ee
Constellation bug fix
jsuarez5341 Feb 3, 2026
096140d
Major bug fix: PPO loss for multidiscrete
jsuarez5341 Feb 3, 2026
706b9e3
merge
jsuarez5341 Feb 3, 2026
c45fad6
Config tweaks
jsuarez5341 Feb 3, 2026
2cb5dd0
Actually fix x axis
jsuarez5341 Feb 4, 2026
aec0826
Merge branch '4.0' of https://github.com/pufferai/pufferlib into 4.0
jsuarez5341 Feb 4, 2026
5df75a2
Stable breakout sweep commit
jsuarez5341 Feb 5, 2026
623c3be
Stable breakout sweep commit
jsuarez5341 Feb 5, 2026
7cd2ce7
Merge sweeps config
jsuarez5341 Feb 5, 2026
8b4594c
Initial multigpu
jsuarez5341 Feb 5, 2026
516a635
Continuous action space support. Drone does not train discrete or con…
jsuarez5341 Feb 5, 2026
e31c6dc
Continuous squared test env
jsuarez5341 Feb 5, 2026
4afd0e5
Sweep fixes
jsuarez5341 Feb 6, 2026
dad389d
Delete old unused dev code
jsuarez5341 Feb 6, 2026
4c112f0
dispatch macro + templated kernel
jsuarez5341 Feb 7, 2026
f387ca1
Static precision compilation = no templating or dispatch
jsuarez5341 Feb 7, 2026
2b83451
Move modules to .cu = no launch fn wrappers
jsuarez5341 Feb 7, 2026
44b8ec3
Initial 2048 sweep. Need to bump up total timesteps. 17k score
jsuarez5341 Feb 7, 2026
e359766
Initial policy cleanup
jsuarez5341 Feb 7, 2026
07ccaea
Un-nest graphing
jsuarez5341 Feb 7, 2026
c50eaaa
Refactor pufferlib.cpp
jsuarez5341 Feb 7, 2026
9de98c4
ocean cpp
jsuarez5341 Feb 7, 2026
00c391a
clean commit
jsuarez5341 Feb 8, 2026
2529dd8
add modules.h
jsuarez5341 Feb 8, 2026
f2069a7
exprimental for fun project
jsuarez5341 Feb 8, 2026
71bdf65
Basic pathing works
jsuarez5341 Feb 8, 2026
56a735d
Profile kernels cleanup
jsuarez5341 Feb 9, 2026
42fc700
merge
jsuarez5341 Feb 9, 2026
30796f2
Delete old env bindings - moved to ocean
jsuarez5341 Feb 9, 2026
0837e76
delete old tests
jsuarez5341 Feb 9, 2026
abc8bb0
Add benchmark env
jsuarez5341 Feb 9, 2026
2385af5
benchmark bind
jsuarez5341 Feb 9, 2026
964067f
benchmark ini
jsuarez5341 Feb 9, 2026
3d67cdb
bandwidth test
jsuarez5341 Feb 9, 2026
d70e557
Env vec speeds look reasonable in testing
jsuarez5341 Feb 9, 2026
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 22 additions & 0 deletions assets/dash.css
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
:root {
--font-color: #F1F1F1;
--dropdown-bg: #005050;
}

body {
background-color: black !important;
color: var(--font-color) !important;
}

.rc-slider-mark-text {
color: var(--font-color) !important;
}

.Select-control, .Select-menu-outer, .Select-value-label, .Select-option {
color: var(--font-color) !important;
background-color: var(--dropdown-bg) !important;
}

h1, h2, h3, h4, h5, h6 {
color: var(--font-color) !important;
}
195 changes: 195 additions & 0 deletions cache_data.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,195 @@
import numpy as np

import json
import glob
import os


env_names = sorted([
'breakout',
#'impulse_wars',
#'pacman',
#'tetris',
#'g2048',
#'moba',
#'pong',
#'tower_climb',
#'grid',
#'nmmo3',
#'snake',
#'tripletriad'
])

HYPERS = [
'train/learning_rate',
'train/ent_coef',
'train/gamma',
'train/gae_lambda',
'train/vtrace_rho_clip',
'train/vtrace_c_clip',
'train/clip_coef',
'train/vf_clip_coef',
'train/vf_coef',
'train/max_grad_norm',
'train/beta1',
'train/beta2',
'train/eps',
'train/prio_alpha',
'train/prio_beta0',
'train/horizon',
'train/replay_ratio',
'train/minibatch_size',
'policy/hidden_size',
'vec/total_agents',
]

ALL_KEYS = [
'agent_steps',
'cost',
'environment/score',
'environment/perf'
] + HYPERS

def pareto_idx(steps, costs, scores):
idxs = []
for i in range(len(steps)):
better = [scores[j] >= scores[i] and
costs[j] < costs[i] and steps[j] < steps[i]
for j in range(len(scores))]
if not any(better):
idxs.append(i)

return idxs

def load_sweep_data(path):
data = {}
keys = None
for fpath in glob.glob(path):
if 'cache.json' in fpath:
continue

with open(fpath, 'r') as f:
exp = json.load(f)

if not data:
for kk in exp.keys():
if kk == 'data':
for k, v in exp[kk][-1].items():
data[k] = []
else:
data[kk] = []

discard = False
for kk in list(data.keys()):
if kk not in exp and kk not in exp['data'][-1]:
discard = True
break

if discard:
continue

for kk in list(data.keys()):
if kk in exp:
v = exp[kk]
sweep_key = f'sweep/{kk}/distribution'
if sweep_key in data and exp[sweep_key] == 'logit_normal':
v = 1 - v
elif kk in ('train/vtrace_rho_clip', 'train/vtrace_c_clip'):
v = max(v, 0.1)

data[kk].append(v)
else:
data[kk].append(exp['data'][-1][kk])

steps = data['agent_steps']
costs = data['cost']
scores = data['environment/score']

idxs = pareto_idx(steps, costs, scores)

# Filter to pareto
for k in data:
data[k] = [data[k][i] for i in idxs]

# Monkey patch: Cap performance
data['environment/perf'] = [min(e, 1.0) for e in data['environment/perf']]

# Monkey patch: Adjust steps by frameskip if present
if 'env/frameskip' in data:
skip = data['env/frameskip']
data['agent_steps'] = [n*m for n, m in zip(data['agent_steps'], skip)]

return data

def cached_sweep_load(path, env_name):
cache_file = os.path.join(path, 'c_cache.json')
if not os.path.exists(cache_file):
data = load_sweep_data(os.path.join(path, '*.json'))
with open(cache_file, 'w') as f:
json.dump(data, f)

with open(cache_file, 'r') as f:
data = json.load(f)

print(f'Loaded {env_name}')
return data

def compute_tsne():
data = {name: cached_sweep_load(f'experiments/logs/puffer_{name}', name) for name in env_names}

flat = []
flat_mmin = []
flat_mmax = []
for env in env_names:
flat.append(np.stack([data[env][hyper] for hyper in HYPERS], axis=1))
flat_mmin.append(np.stack([data[env][f'sweep/{hyper}/min'] for hyper in HYPERS], axis=1))
flat_mmax.append(np.stack([data[env][f'sweep/{hyper}/max'] for hyper in HYPERS], axis=1))

flat_distribution = [data[env][f'sweep/{hyper}/distribution'] for env in env_names for hyper in HYPERS]

flat = np.concatenate(flat, axis=0)
flat_mmin = np.concatenate(flat_mmin, axis=0).min(axis=0)
flat_mmax = np.concatenate(flat_mmax, axis=0).max(axis=0)

normed = flat.copy()
for i in range(len(HYPERS)):
dist = flat_distribution[i]
if 'log' in dist or 'pow2' in dist:
flat_mmin[i] = np.log(flat_mmin[i])
flat_mmax[i] = np.log(flat_mmax[i])
normed[:, i] = np.log(flat[:, i])

normed[:, i] = (normed[:, i] - flat_mmin[i]) / (flat_mmax[i] - flat_mmin[i])

from sklearn.manifold import TSNE
proj = TSNE(n_components=2)
reduced = None
try:
reduced = proj.fit_transform(normed)
except ValueError:
print('Warning: TSNE failed. Skipping TSNE')

row = 0
for env in env_names:
'''
for i, hyper in enumerate(HYPERS):
sz = len(data[env][hyper])
data[env][hyper] = normed[row:row+sz, i].tolist()
'''
sz = len(data[env]['agent_steps'])

data[env] = {k: v for k, v in data[env].items() if k in ALL_KEYS}
if reduced is not None:
data[env]['tsne1'] = reduced[row:row+sz, 0].tolist()
data[env]['tsne2'] = reduced[row:row+sz, 1].tolist()
else:
data[env]['tsne1'] = np.random.rand(sz).tolist()
data[env]['tsne2'] = np.random.rand(sz).tolist()

row += sz
print(f'Env {env} has {sz} points')

json.dump(data, open('all_cache.json', 'w'))

if __name__ == '__main__':
compute_tsne()
Loading
Loading