Rename to EpiWM: model card, analysis columns (epiwm_*), animations

README.md

@@ -8,9 +8,9 @@ tags:
 - lewm
 ---
 
-# EpiJEPA-LeWM: LeWorldModel trained with epiplexity instead of SIGReg
+# EpiWM: LeWorldModel trained with epiplexity instead of SIGReg
 
-These are the best checkpoints from replacing [LeWorldModel](https://github.com/lucas-maes/le-wm)'s SIGReg
+EpiWM: these are the best checkpoints from replacing [LeWorldModel](https://github.com/lucas-maes/le-wm)'s SIGReg
 anti-collapse term with the epiplexity score from EpiJEPA. The code, the full analysis and our replication of LeWM's
 SIGReg recipe are in [the-puzzler/epijepa](https://github.com/the-puzzler/epijepa) (`worldmodel/`).
 
@@ -41,14 +41,14 @@ where S is the epiplexity of the embeddings with respect to a frozen random CNN
 
 ```bash
 # with the repo's worldmodel/ folder on PYTHONPATH (config.json refers to module.ProjectorBN)
-hf download basilboy/epijepa-lewm --local-dir $STABLEWM_HOME/checkpoints/epijepa_lewm
+hf download basilboy/epiwm --local-dir $STABLEWM_HOME/checkpoints/epiwm
 cd epijepa/worldmodel
-python eval.py --config-name=pusht.yaml policy=epijepa_lewm/pusht        # LeWM's eval: paper-50
+python eval.py --config-name=pusht.yaml policy=epiwm/pusht        # LeWM's eval: paper-50
 ```
 
 ```python
 import stable_worldmodel as swm
-model = swm.wm.utils.load_pretrained("epijepa_lewm/pusht")   # resolved under $STABLEWM_HOME/checkpoints
+model = swm.wm.utils.load_pretrained("epiwm/pusht")   # resolved under $STABLEWM_HOME/checkpoints
 ```
 
 The weights are plain state dicts (`weights.pt`, fp32, about 72 MB), saved with stable-worldmodel's `save_pretrained`
@@ -61,7 +61,7 @@ folder's README, and the scripts there regenerate everything.
 
 | Path | What it is |
 |---|---|
-| `embeddings/embeddings_<env>.csv` | 3000 random frames per environment: episode, step, true state, PCA and t-SNE coordinates for EpiJEPA and released LeWM |
+| `embeddings/embeddings_<env>.csv` | 3000 random frames per environment: episode, step, true state, PCA and t-SNE coordinates for EpiWM and released LeWM |
 | `embeddings/trajectories_<env>.csv` + `embeddings/videos/<env>_ep<N>.mp4` | one full episode per environment. Row `step == i` is frame `i` of the video, with its true state and coordinates in both models' PCA spaces |
 | `embeddings/animations/<env>_ep<N>_pca.mp4` | the episode video side by side with the moving point in both PCA spaces |
 | `embeddings/pca_basis_<env>.npz`, `pca_variance_<env>.csv` | the PCA bases (mean, 50×192 components) and explained variance |

analysis/embeddings/animations/cube_ep348_pca.mp4

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:54e05421df7f8272f825955f36f6cd319963666dbd53c55c86d4870bfb38ec4b
-size 346210
+oid sha256:33a5c3cfb133b968fce88fe2ef5cf42748b8582888324cf20d74ba1137a43ead
+size 346878

analysis/embeddings/animations/pusht_ep630_pca.mp4

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7537fe65b7e907d88f261ac8a996187578919aa1eddf769caa299114d151ee3d
-size 137592
+oid sha256:a248dcd8d8d908abbcf3d2d81a24f4c16cb1019e569f6869c9e7f9e7274256d0
+size 137259

analysis/embeddings/animations/reacher_ep348_pca.mp4

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:036d8fa8fcc5ea83f9a8ca920112a2a350624687884e3fc4f08ba5f50a0ef6e9
-size 166438
+oid sha256:9cc6188af06383e26db428dc1622e8bf1a81c05e14040be90926bc79ad537e64
+size 166852

analysis/embeddings/embeddings_cube.csv

@@ -1,4 +1,4 @@
-row,block_x,block_y,block_z,block_yaw,effector_x,effector_y,effector_z,gripper_opening,episode,step,epijepa_pca1,epijepa_pca2,epijepa_tsne1,epijepa_tsne2,released_lewm_pca1,released_lewm_pca2,released_lewm_tsne1,released_lewm_tsne2
+row,block_x,block_y,block_z,block_yaw,effector_x,effector_y,effector_z,gripper_opening,episode,step,epiwm_pca1,epiwm_pca2,epiwm_tsne1,epiwm_tsne2,released_lewm_pca1,released_lewm_pca2,released_lewm_tsne1,released_lewm_tsne2
 381,0.43828838741161247,-0.17220837913388948,0.019960724713120808,-2.1225675675148956,0.3601970885766833,-0.022062517275799644,0.24379003370469698,0.07444620500412498,1,180,0.8352535573718798,4.979322572883428,6.0977063,34.056103,1.0031014141823358,-1.4827638048959317,-22.157936,-27.146154
 445,0.39662011748401604,0.16065298450256715,0.27116076580315074,-0.6706181919201888,0.3923416583503355,0.1558926658213467,0.26679717122410457,0.5614144710592923,2,43,-6.796934420894286,0.5731233102416263,-46.854904,23.35277,-0.15558196146452088,0.025691896110798945,-0.73125976,17.958942
 603,0.3305884269102304,0.16561672326769583,0.02,-0.7338876804955067,0.42496850313414825,-5.474531403616736e-05,0.2749649648972977,0.0,3,0,-1.7689984349298336,14.745612978899857,-4.309732,55.950237,1.7670342521375393,-0.6817517678027578,23.94193,32.52116

analysis/embeddings/embeddings_pusht.csv

@@ -1,4 +1,4 @@
-row,agent_x,agent_y,block_x,block_y,block_angle,episode,step,epijepa_pca1,epijepa_pca2,epijepa_tsne1,epijepa_tsne2,released_lewm_pca1,released_lewm_pca2,released_lewm_tsne1,released_lewm_tsne2
+row,agent_x,agent_y,block_x,block_y,block_angle,episode,step,epiwm_pca1,epiwm_pca2,epiwm_tsne1,epiwm_tsne2,released_lewm_pca1,released_lewm_pca2,released_lewm_tsne1,released_lewm_tsne2
 443,368.93164,173.74062,192.67061,335.15677,2.9476998,4,7,-2.132587934270504,4.915789494720234,-11.789562,15.191289,-1.8495206377324376,1.701001597269687,34.931675,12.7589855
 517,159.5481,430.72977,192.09763,321.4666,0.8671616,4,81,1.1696597014312244,1.207833060143149,26.691744,35.578964,-0.9511822749297324,0.7191298168117688,26.07366,37.464878
 701,84.056694,302.50723,192.32678,335.4081,1.6003228,6,47,1.2471565406973077,1.5224858111893187,37.98232,22.114403,-0.7341384125912208,-1.0499743024756831,33.32419,52.916443

analysis/embeddings/embeddings_reacher.csv

@@ -1,4 +1,4 @@
-row,finger_x,finger_y,joint_0,joint_1,episode,step,epijepa_pca1,epijepa_pca2,epijepa_tsne1,epijepa_tsne2,released_lewm_pca1,released_lewm_pca2,released_lewm_tsne1,released_lewm_tsne2
+row,finger_x,finger_y,joint_0,joint_1,episode,step,epiwm_pca1,epiwm_pca2,epiwm_tsne1,epiwm_tsne2,released_lewm_pca1,released_lewm_pca2,released_lewm_tsne1,released_lewm_tsne2
 381,-0.10705846469072372,0.19193140516450008,-3.790058959817336,-0.8270430514210889,1,180,-1.2668852149375376,0.309783137118494,-24.175495,22.509869,-1.5320085757249573,-0.21906824817158338,-1.263843,-15.103656
 445,0.026021927551123314,0.09981630270642164,0.18925158679429793,2.2530456592487913,2,43,0.6583230042091309,-0.620569646045575,56.43641,-4.145483,-0.11939089003216766,-0.3930117848875966,30.095078,21.549139
 603,-0.03772492262847432,-0.2032013123896669,-1.2210702643588256,-1.0665781899490268,3,0,0.6755494104735569,1.5543090962452968,-27.279623,7.745932,-1.4572931477975044,0.8560067130359398,-27.859522,-3.269776

analysis/embeddings/embeddings_tworoom.csv

@@ -1,4 +1,4 @@
-row,agent_x,agent_y,episode,step,epijepa_pca1,epijepa_pca2,epijepa_tsne1,epijepa_tsne2,released_lewm_pca1,released_lewm_pca2,released_lewm_tsne1,released_lewm_tsne2
+row,agent_x,agent_y,episode,step,epiwm_pca1,epiwm_pca2,epiwm_tsne1,epiwm_tsne2,released_lewm_pca1,released_lewm_pca2,released_lewm_tsne1,released_lewm_tsne2
 174,94.344894,136.12839,2,35,-5.88999296600558,-8.388116136965117,3.9403484,-57.028877,5.853797697209181,-1.055108438369035,45.59889,28.819199
 203,99.5,62.80535,2,64,-6.37835678183774,1.9414454309521256,-13.571716,4.913951,1.3942764015642026,0.043212414089021656,-7.0675254,60.96116
 276,108.89513,39.287823,3,56,-7.439050965436256,2.817969790710875,-52.870106,10.281317,-0.4990495745683413,0.12150834119128173,-53.00016,5.411207

analysis/embeddings/pca_basis_cube.npz

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4860d2fdb0caf19f90c0020c0f51870dda7e91eb90c097fd0979afe42391fc03
-size 159170
+oid sha256:129018af792f95f7e1918d53af962f8cd89f4b80e0ffe9bc76110ca080b31e61
+size 159158

analysis/embeddings/pca_basis_pusht.npz

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:933c1919c4b23157b0aeff8bb5ace1121d565186d857c4b7ab3720fa87f7b35f
-size 159170
+oid sha256:f8479c2b1a296783c46d3c6b1bbf395d5bb9785ebd95ec02e9ece3ce88623ece
+size 159158

analysis/embeddings/pca_basis_reacher.npz

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:80a295f6d9da813432a2d2c988fd7fbb5a8bd222f04d2fe0dae0380b360b0bc7
-size 159170
+oid sha256:17142b667fc57df4af58ec6959ca6b480d3b00ba87d30f269cb0b1f6f8070bdf
+size 159158

analysis/embeddings/pca_basis_tworoom.npz

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e9f12754105ca776cafd9f3639c2b2950d563de6667e7d6605e63e2a2ef177ec
-size 159170
+oid sha256:f36de69cb7599a6964798a6413047f6d1f34c3b18100ee8fea7096ffa0cb63f6
+size 159158

analysis/embeddings/pca_variance_cube.csv

@@ -1,4 +1,4 @@
-pc,epijepa,released_lewm
+pc,epiwm,released_lewm
 1,0.17887712709172754,0.015335943358594988
 2,0.04871422592612006,0.015094634654122979
 3,0.03776308129531492,0.014847141790040218

analysis/embeddings/pca_variance_pusht.csv

@@ -1,4 +1,4 @@
-pc,epijepa,released_lewm
+pc,epiwm,released_lewm
 1,0.019655302478823803,0.01953987770001811
 2,0.017791780369069355,0.019269503909159137
 3,0.015809544524634872,0.018365139710839987

analysis/embeddings/pca_variance_reacher.csv

@@ -1,4 +1,4 @@
-pc,epijepa,released_lewm
+pc,epiwm,released_lewm
 1,0.013580246209229334,0.020892089389225325
 2,0.013433666049261552,0.020744862413771548
 3,0.013094447113958033,0.020011791835693443

analysis/embeddings/pca_variance_tworoom.csv

@@ -1,4 +1,4 @@
-pc,epijepa,released_lewm
+pc,epiwm,released_lewm
 1,0.7621223855507904,0.025556580266749365
 2,0.04713827504795719,0.022618734492950602
 3,0.02774698014299456,0.020938351555675246

analysis/embeddings/trajectories_cube.csv

@@ -1,4 +1,4 @@
-block_x,block_y,block_z,block_yaw,effector_x,effector_y,effector_z,gripper_opening,episode,step,epijepa_pca1,epijepa_pca2,epijepa_pca3,released_lewm_pca1,released_lewm_pca2,released_lewm_pca3
+block_x,block_y,block_z,block_yaw,effector_x,effector_y,effector_z,gripper_opening,episode,step,epiwm_pca1,epiwm_pca2,epiwm_pca3,released_lewm_pca1,released_lewm_pca2,released_lewm_pca3
 0.48854200432386397,-0.10413173609505669,0.02,0.4308283522953331,0.42499995349929687,-5.927090471467913e-07,0.2749988821824685,0.0,348,0,1.053630976223653,14.670953678831594,-12.202479508898744,0.7417438299819984,2.101637969015291,-0.269052059965243
 0.48854200432390543,-0.10413173609498974,0.01994887806559409,0.4308283522953331,0.42625019000939307,0.003943601487563314,0.27476790139514695,0.0014282461100838587,348,1,0.8179763491237622,13.692856482245801,-10.35321811219835,0.7606178297027211,2.0791247417817074,-0.2641208946710431
 0.48854200432391115,-0.10413173609498036,0.01995989911734845,0.4308283522953331,0.43494576916749184,-0.0028163646564977503,0.2686940223202149,0.0022814217956405587,348,2,0.7215075431598577,9.017739563639525,-5.042794275001991,0.7973061476916432,2.036395414440026,-0.2795850533749603

analysis/embeddings/trajectories_pusht.csv

@@ -1,4 +1,4 @@
-agent_x,agent_y,block_x,block_y,block_angle,episode,step,epijepa_pca1,epijepa_pca2,epijepa_pca3,released_lewm_pca1,released_lewm_pca2,released_lewm_pca3
+agent_x,agent_y,block_x,block_y,block_angle,episode,step,epiwm_pca1,epiwm_pca2,epiwm_pca3,released_lewm_pca1,released_lewm_pca2,released_lewm_pca3
 359.0,368.0,333.69165,212.01947,5.6028233,630,0,-3.3272791421014274,-1.4326715571161772,-2.459223440078408,-1.4669347822006467,0.2744595501349158,-0.18652049914651916
 351.13025,367.04202,333.69165,212.01947,5.6028233,630,1,-3.425589837161334,-1.3769855437616911,-2.719298392725144,-1.499142102319607,0.28306750074962317,-0.17932135326561377
 343.6903,366.11856,333.69165,212.01947,5.6028233,630,2,-3.5120688074093818,-1.3196211297234828,-2.932684411654445,-1.6533357586954323,0.27371211059090256,-0.07585147887875276

analysis/embeddings/trajectories_reacher.csv

@@ -1,4 +1,4 @@
-finger_x,finger_y,joint_0,joint_1,episode,step,epijepa_pca1,epijepa_pca2,epijepa_pca3,released_lewm_pca1,released_lewm_pca2,released_lewm_pca3
+finger_x,finger_y,joint_0,joint_1,episode,step,epiwm_pca1,epiwm_pca2,epiwm_pca3,released_lewm_pca1,released_lewm_pca2,released_lewm_pca3
 -0.2093465604903495,0.068677481472573,2.4168500351601585,0.8154975755690219,348,0,0.11704597106530587,-0.20303088468701744,0.19242426959346265,0.3617641430781829,-0.5103230850227587,-0.9582022349876036
 -0.21213527252266062,0.07909407548766087,2.4464579960344,0.6765051422126989,348,1,-0.04761594700084781,-0.16227941403693333,0.15353663515314567,0.46896025077516296,0.031358053214598926,-0.7328210759201527
 -0.2129968997293278,0.08026541872759721,2.458599827076246,0.6452227021348722,348,2,-0.1337382540829443,-0.10490819461442207,0.12148819281926224,0.4991737928409297,0.1455504751053702,-0.6659664868607258

analysis/embeddings/trajectories_tworoom.csv

@@ -1,4 +1,4 @@
-agent_x,agent_y,episode,step,epijepa_pca1,epijepa_pca2,epijepa_pca3,released_lewm_pca1,released_lewm_pca2,released_lewm_pca3
+agent_x,agent_y,episode,step,epiwm_pca1,epiwm_pca2,epiwm_pca3,released_lewm_pca1,released_lewm_pca2,released_lewm_pca3
 71.12141,104.50593,4748,0,-6.102547536870161,-4.415961715074084,1.571122984699361,1.2559883996311252,-1.3604212302404646,-3.689166183427508
 66.12141,106.34553,4748,1,-6.311987011603993,-4.928670632123013,1.3489407849805974,-0.28882979657483593,-4.730573108524823,-2.873377336799558
 66.22549,111.34553,4748,2,-6.544297064527109,-5.731511646968797,1.0731443198171375,-0.11409361177164373,-4.299157252391719,-1.7895383806781389

analysis/probes/probes.json

@@ -1,6 +1,6 @@
 {
  "tworoom": {
-  "epijepa": {
+  "epiwm": {
    "checkpoint": "tw_epi_0.03/weights_step30000.pt",
    "linear_r2": {
     "agent_x": 0.9974128927299404,
@@ -62,7 +62,7 @@
   }
  },
  "pusht": {
-  "epijepa": {
+  "epiwm": {
    "checkpoint": "pusht_epi_0.1_60k_seed3/weights_step60000.pt",
    "linear_r2": {
     "agent_x": 0.9639334393958146,
@@ -140,7 +140,7 @@
   }
  },
  "cube": {
-  "epijepa": {
+  "epiwm": {
    "checkpoint": "cube_epi_0.03_seed3/weights_step60000.pt",
    "linear_r2": {
     "block_x": 0.9385609012362691,
@@ -230,7 +230,7 @@
   }
  },
  "reacher": {
-  "epijepa": {
+  "epiwm": {
    "checkpoint": "rch_epi_0.3_200k/weights_step200000.pt",
    "linear_r2": {
     "finger_x": 0.9992451291317158,