Updated README

README.md

@@ -13,6 +13,7 @@ We share models trained for one seed (123), whereas the paper contained weights
 ### Usage
 
 ```
+git clone https://huggingface.co/edbeeching/decision_transformer_atari
 conda env create -f conda_env.yml
 ```
 
@@ -20,7 +21,25 @@ Then, you can use the model like this:
 
 ```python
 
-
-
-
+from decision_transform_atari import GPTConfig, GPT
+
+vocab_size = 4
+block_size = 90
+model_type = "reward_conditioned"
+timesteps = 2654
+
+mconf = GPTConfig(
+    vocab_size,
+    block_size,
+    n_layer=6,
+    n_head=8,
+    n_embd=128,
+    model_type=model_type,
+    max_timestep=timesteps,
+)
+model = GPT(mconf)
+
+checkpoint_path = "checkpoints/Breakout_123.pth"  # or Pong, Qbert, Seaquest
+checkpoint = torch.load(checkpoint_path)
+model.load_state_dict(checkpoint)
 ```