diff --git a/.idea/.gitignore b/.idea/.gitignore
new file mode 100644
index 00000000..13566b81
--- /dev/null
+++ b/.idea/.gitignore
@@ -0,0 +1,8 @@
+# Default ignored files
+/shelf/
+/workspace.xml
+# Editor-based HTTP Client requests
+/httpRequests/
+# Datasource local storage ignored files
+/dataSources/
+/dataSources.local.xml
diff --git a/.idea/humanoid-bench.iml b/.idea/humanoid-bench.iml
new file mode 100644
index 00000000..2f9173b4
--- /dev/null
+++ b/.idea/humanoid-bench.iml
@@ -0,0 +1,19 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.idea/inspectionProfiles/Project_Default.xml b/.idea/inspectionProfiles/Project_Default.xml
new file mode 100644
index 00000000..ac21435f
--- /dev/null
+++ b/.idea/inspectionProfiles/Project_Default.xml
@@ -0,0 +1,6 @@
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.idea/inspectionProfiles/profiles_settings.xml b/.idea/inspectionProfiles/profiles_settings.xml
new file mode 100644
index 00000000..105ce2da
--- /dev/null
+++ b/.idea/inspectionProfiles/profiles_settings.xml
@@ -0,0 +1,6 @@
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.idea/misc.xml b/.idea/misc.xml
new file mode 100644
index 00000000..778f9db7
--- /dev/null
+++ b/.idea/misc.xml
@@ -0,0 +1,7 @@
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.idea/modules.xml b/.idea/modules.xml
new file mode 100644
index 00000000..3c7ac11c
--- /dev/null
+++ b/.idea/modules.xml
@@ -0,0 +1,8 @@
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.idea/vcs.xml b/.idea/vcs.xml
new file mode 100644
index 00000000..35eb1ddf
--- /dev/null
+++ b/.idea/vcs.xml
@@ -0,0 +1,6 @@
+
+
+
+
+
+
\ No newline at end of file
diff --git a/commands b/commands
new file mode 100644
index 00000000..c648cd5e
--- /dev/null
+++ b/commands
@@ -0,0 +1,7 @@
+conda create -n humanoidbench python=3.11
+conda activate humanoidbench
+pip install -e .
+pip install "jax[cuda12]==0.4.28"
+pip install -r requirements_jaxrl.txt
+pip install -r requirements_dreamer.txt
+pip install -r requirements_tdmpc.txt
diff --git a/pip_packages.txt b/pip_packages.txt
new file mode 100644
index 00000000..562449c3
--- /dev/null
+++ b/pip_packages.txt
@@ -0,0 +1,180 @@
+Package Version Editable project location
+---------------------------- ------------ ----------------------------------
+absl-py 2.2.1
+antlr4-python3-runtime 4.9.3
+asttokens 3.0.0
+astunparse 1.6.3
+attrs 25.3.0
+blinker 1.9.0
+brax 0.9.4
+certifi 2025.1.31
+charset-normalizer 3.4.1
+chex 0.1.89
+click 8.1.8
+cloudpickle 3.0.0
+colored 2.2.4
+contextlib2 21.6.0
+contourpy 1.3.1
+cycler 0.12.1
+decorator 4.4.2
+distrax 0.1.5
+dm_control 1.0.20
+dm-env 1.6
+dm-sonnet 2.0.2
+dm-tree 0.1.9
+docker-pycreds 0.4.0
+embodied 0.0.0 /home/skr/humanoid-bench/dreamerv3
+etils 1.12.2
+executing 2.2.0
+Farama-Notifications 0.0.4
+filelock 3.18.0
+Flask 3.1.0
+flask-cors 5.0.1
+flatbuffers 25.2.10
+flax 0.8.5
+fonttools 4.56.0
+fsspec 2025.3.0
+gast 0.6.0
+gitdb 4.0.12
+GitPython 3.1.44
+glfw 2.8.0
+google-pasta 0.2.0
+grpcio 1.71.0
+gym 0.26.2
+gym-notices 0.0.8
+gymnasium 0.29.1
+gymnax 0.0.8
+h5py 3.13.0
+humanize 4.12.2
+humanoid_bench 0.2 /home/skr/humanoid-bench
+hydra-core 1.3.2
+hydra-submitit-launcher 1.2.0
+idna 3.10
+imageio 2.34.2
+imageio-ffmpeg 0.6.0
+importlib_resources 6.5.2
+ipdb 0.13.13
+ipython 9.0.2
+ipython_pygments_lexers 1.1.1
+itsdangerous 2.2.0
+jax 0.5.3
+jax-cuda12-pjrt 0.5.3
+jax-cuda12-plugin 0.5.3
+jaxlib 0.5.3
+jaxopt 0.8.3
+jaxrl_m 0.1.0 /home/skr/humanoid-bench/jaxrl_m
+jedi 0.19.2
+Jinja2 3.1.6
+keras 3.9.1
+kiwisolver 1.4.8
+labmaze 1.0.6
+libclang 18.1.1
+lxml 5.3.1
+Markdown 3.7
+markdown-it-py 3.0.0
+MarkupSafe 3.0.2
+matplotlib 3.9.0
+matplotlib-inline 0.1.7
+mdurl 0.1.2
+ml_collections 0.1.1
+ml_dtypes 0.5.1
+moviepy 1.0.3
+mpmath 1.3.0
+msgpack 1.0.8
+mujoco 3.1.6
+mujoco-mjx 3.1.6
+multiagent 0.0.1
+namex 0.0.8
+natsort 8.4.0
+nest-asyncio 1.6.0
+networkx 3.4.2
+numpy 1.26.4
+nvidia-cublas-cu12 12.1.3.1
+nvidia-cuda-cupti-cu12 12.1.105
+nvidia-cuda-nvcc-cu12 12.8.93
+nvidia-cuda-nvrtc-cu12 12.1.105
+nvidia-cuda-runtime-cu12 12.1.105
+nvidia-cudnn-cu12 8.9.2.26
+nvidia-cufft-cu12 11.0.2.54
+nvidia-curand-cu12 10.3.2.106
+nvidia-cusolver-cu12 11.4.5.107
+nvidia-cusparse-cu12 12.1.0.106
+nvidia-nccl-cu12 2.20.5
+nvidia-nvjitlink-cu12 12.8.93
+nvidia-nvtx-cu12 12.1.105
+omegaconf 2.3.0
+opencv-python 4.10.0.84
+opt_einsum 3.4.0
+optax 0.2.2
+optree 0.14.1
+orbax-checkpoint 0.11.5
+packaging 24.2
+pandas 2.2.2
+parso 0.8.4
+pexpect 4.9.0
+pillow 11.1.0
+pip 25.0
+platformdirs 4.3.7
+plotly 5.22.0
+proglog 0.1.10
+prompt_toolkit 3.0.50
+protobuf 4.25.6
+psutil 7.0.0
+ptyprocess 0.7.0
+pure_eval 0.2.3
+Pygments 2.19.1
+PyOpenGL 3.1.9
+pyparsing 3.2.3
+pyquaternion 0.9.9
+python-dateutil 2.9.0.post0
+pytinyrenderer 0.0.14
+pytz 2025.2
+PyYAML 6.0.2
+pyzmq 26.0.3
+requests 2.32.3
+rich 13.7.1
+ruamel.yaml 0.18.6
+ruamel.yaml.clib 0.2.12
+scipy 1.15.2
+seaborn 0.13.2
+sentry-sdk 2.24.1
+setproctitle 1.3.5
+setuptools 75.8.0
+simplejson 3.20.1
+six 1.17.0
+smmap 5.0.2
+stack-data 0.6.3
+submitit 1.5.2
+sympy 1.13.3
+tabulate 0.9.0
+tdmpc2 0.0 /home/skr/humanoid-bench/tdmpc2
+tenacity 9.0.0
+tensorboard 2.16.2
+tensorboard-data-server 0.7.2
+tensorboardX 2.6.2.2
+tensordict 0.4.0
+tensorflow 2.16.2
+tensorflow-io-gcs-filesystem 0.37.1
+tensorflow-probability 0.24.0
+tensorstore 0.1.73
+termcolor 2.4.0
+tf_keras 2.16.0
+toolz 1.0.0
+torch 2.3.1+cu121
+torchaudio 2.3.1+cu121
+torchrl 0.4.0
+torchvision 0.18.1+cu121
+tqdm 4.66.4
+traitlets 5.14.3
+treescope 0.1.9
+trimesh 4.6.6
+triton 2.3.1
+typing_extensions 4.13.0
+tzdata 2025.2
+urllib3 2.3.0
+wandb 0.17.3
+wcwidth 0.2.13
+Werkzeug 3.1.3
+wheel 0.45.1
+wrapt 1.17.2
+zipp 3.21.0
diff --git a/requirements_jaxrl.txt b/requirements_jaxrl.txt
index fc649525..b43a9de3 100644
--- a/requirements_jaxrl.txt
+++ b/requirements_jaxrl.txt
@@ -2,4 +2,7 @@
ml_collections==0.1.1
flax==0.8.5
distrax==0.1.5
-tf_keras==2.16.0
\ No newline at end of file
+tf_keras==2.16.0
+
+
+#pip uninstall -y jax jaxlib jaxopt jaxrl_m jax-cuda12-plugin jax-cuda12-pjrt
diff --git a/requirements_tdmpc.txt b/requirements_tdmpc.txt
index 6ac36e1d..c4054890 100644
--- a/requirements_tdmpc.txt
+++ b/requirements_tdmpc.txt
@@ -1,11 +1,15 @@
--editable tdmpc2
torch==2.3.1
-torchaudio==2.3.1
-torchrl==0.4.0
torchvision==0.18.1
+torchaudio==2.3.1
+--extra-index-url https://download.pytorch.org/whl/cu121
+
hydra-core==1.3.2
hydra-submitit-launcher==1.2.0
pyquaternion==0.9.9
tensordict==0.4.0
pandas==2.2.2
termcolor==2.4.0
+jax[cuda12]==0.4.28
+torchrl==0.4.0
+
diff --git a/setup.py b/setup.py
index c93bee9d..55a5df87 100644
--- a/setup.py
+++ b/setup.py
@@ -15,7 +15,7 @@
"imageio==2.34.2",
"gymnax==0.0.8",
"brax==0.9.4",
- "torch==2.3.1",
+ #"torch==2.3.1",
"opencv-python==4.10.0.84",
"natsort==8.4.0",
]
diff --git a/test_env_img.png b/test_env_img.png
index 6f2031f5..c64501bf 100644
Binary files a/test_env_img.png and b/test_env_img.png differ