Update README.md

README.md

@@ -17,7 +17,7 @@ This model is trained on the Jumanji snake environment
 
 ### How to use
 
-Go to the jumanji repo for the primary model and requirements.
+Go to the jumanji repo for the primary model and requirements. Clone the repo and navigate to the root directory.
 
 ```
 pip install --quiet -U pip -r ../requirements/requirements-train.txt ../.
@@ -26,7 +26,6 @@ pip install --quiet -U pip -r ../requirements/requirements-train.txt ../.
 Below is an example script for loading and running the Jumanji model
 
 ```python
-
 import pickle
 import joblib
 
@@ -39,17 +38,16 @@ from jumanji.training.setup_train import setup_agent, setup_env
 from jumanji.training.utils import first_from_device
 
 # initialise the config
-with initialize(version_base=None, config_path="../jumanji/training/configs"):
-    cfg = compose(config_name="config.yaml", overrides=["env=connector", "agent=a2c"])
+with initialize(version_base=None, config_path="jumanji/training/configs"):
+    cfg = compose(config_name="config.yaml", overrides=["env=snake", "agent=a2c"])
 
 # get model state from HF
-REPO_ID = "YOUR_REPO_ID"
+REPO_ID = "d-byrne/snake-v1_training_state"
 FILENAME = "Snake-v1_training_state"
 
 model_weights = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
 
-# load the model
-with open("training_state","rb") as f:
+with open(model_weights,"rb") as f:
     training_state = pickle.load(f)
 
 params = first_from_device(training_state.params_state.params)