Fix/minedojo (#286)

* fix: multiple envs

* fix: multi-discrete actions

* fix: remove debug prints

* fix: remove debug prints

* fix: removed MINEDOJO_HEADLESS

howto/register_external_algorithm.md

@@ -605,9 +605,9 @@ def ext_sota_main(fabric: Fabric, cfg: Dict[str, Any]):
                     }
                     actions = agent.module(torch_obs)
                     if is_continuous:
-                        real_actions = torch.cat(actions, -1).cpu().numpy()
+                        real_actions = torch.stack(actions, -1).cpu().numpy()
                     else:
-                        real_actions = torch.cat([act.argmax(dim=-1) for act in actions], dim=-1).cpu().numpy()
+                        real_actions = torch.stack([act.argmax(dim=-1) for act in actions], dim=-1).cpu().numpy()
                     actions = torch.cat(actions, -1).cpu().numpy()
 
                 # Single environment step

howto/register_new_algorithm.md

@@ -603,9 +603,9 @@ def sota_main(fabric: Fabric, cfg: Dict[str, Any]):
                     }
                     actions = player.get_actions(torch_obs)
                     if is_continuous:
-                        real_actions = torch.cat(actions, -1).cpu().numpy()
+                        real_actions = torch.stack(actions, -1).cpu().numpy()
                     else:
-                        real_actions = torch.cat([act.argmax(dim=-1) for act in actions], dim=-1).cpu().numpy()
+                        real_actions = torch.stack([act.argmax(dim=-1) for act in actions], dim=-1).cpu().numpy()
                     actions = torch.cat(actions, -1).cpu().numpy()
 
                 # Single environment step

notebooks/dreamer_v3_imagination.ipynb

@@ -222,9 +222,9 @@
     "        real_actions = actions = player.get_actions(preprocessed_obs, mask=mask)\n",
     "        actions = torch.cat(actions, -1).cpu().numpy()\n",
     "        if is_continuous:\n",
-    "            real_actions = torch.cat(real_actions, dim=-1).cpu().numpy()\n",
+    "            real_actions = torch.stack(real_actions, dim=-1).cpu().numpy()\n",
     "        else:\n",
-    "            real_actions = torch.cat([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()\n",
+    "            real_actions = torch.stack([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()\n",
     "\n",
     "    step_data[\"actions\"] = actions.reshape((1, cfg.env.num_envs, -1))\n",
     "    rb_initial.add(step_data, validate_args=cfg.buffer.validate_args)\n",

sheeprl/algos/a2c/a2c.py

@@ -239,9 +239,9 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
                     torch_obs = prepare_obs(fabric, next_obs, num_envs=cfg.env.num_envs)
                     actions, _, values = player(torch_obs)
                     if is_continuous:
-                        real_actions = torch.cat(actions, -1).cpu().numpy()
+                        real_actions = torch.stack(actions, -1).cpu().numpy()
                     else:
-                        real_actions = torch.cat([act.argmax(dim=-1) for act in actions], axis=-1).cpu().numpy()
+                        real_actions = torch.stack([act.argmax(dim=-1) for act in actions], axis=-1).cpu().numpy()
                     actions = torch.cat(actions, -1).cpu().numpy()
 
                     # Single environment step

sheeprl/algos/dreamer_v1/dreamer_v1.py

@@ -581,7 +581,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
                     real_actions = actions = player.get_exploration_actions(torch_obs, mask=mask, step=policy_step)
                     actions = torch.cat(actions, -1).view(cfg.env.num_envs, -1).cpu().numpy()
                     if is_continuous:
-                        real_actions = torch.cat(real_actions, -1).cpu().numpy()
+                        real_actions = torch.stack(real_actions, -1).cpu().numpy()
                     else:
                         real_actions = (
                             torch.cat([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()

sheeprl/algos/dreamer_v2/dreamer_v2.py

@@ -145,7 +145,11 @@ def train(
         # One step of dynamic learning, which take the posterior state, the recurrent state, the action
         # and the observation and compute the next recurrent, prior and posterior states
         recurrent_state, posterior, _, posterior_logits, prior_logits = world_model.rssm.dynamic(
-            posterior, recurrent_state, data["actions"][i : i + 1], embedded_obs[i : i + 1], data["is_first"][i : i + 1]
+            posterior,
+            recurrent_state,
+            data["actions"][i : i + 1],
+            embedded_obs[i : i + 1],
+            data["is_first"][i : i + 1],
         )
         recurrent_states[i] = recurrent_state
         priors_logits[i] = prior_logits
@@ -344,7 +348,10 @@ def train(
     critic_grads = None
     if cfg.algo.critic.clip_gradients is not None and cfg.algo.critic.clip_gradients > 0:
         critic_grads = fabric.clip_gradients(
-            module=critic, optimizer=critic_optimizer, max_norm=cfg.algo.critic.clip_gradients, error_if_nonfinite=False
+            module=critic,
+            optimizer=critic_optimizer,
+            max_norm=cfg.algo.critic.clip_gradients,
+            error_if_nonfinite=False,
         )
     critic_optimizer.step()
 
@@ -606,10 +613,10 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
                     real_actions = actions = player.get_actions(torch_obs, mask=mask)
                     actions = torch.cat(actions, -1).view(cfg.env.num_envs, -1).cpu().numpy()
                     if is_continuous:
-                        real_actions = torch.cat(real_actions, -1).cpu().numpy()
+                        real_actions = torch.stack(real_actions, -1).cpu().numpy()
                     else:
                         real_actions = (
-                            torch.cat([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()
+                            torch.stack([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()
                         )
 
                 step_data["is_first"] = copy.deepcopy(np.logical_or(step_data["terminated"], step_data["truncated"]))

sheeprl/algos/dreamer_v2/utils.py

@@ -150,9 +150,9 @@ def test(
             torch_obs, greedy, {k: v for k, v in torch_obs.items() if k.startswith("mask")}
         )
         if player.actor.is_continuous:
-            real_actions = torch.cat(real_actions, -1).cpu().numpy()
+            real_actions = torch.stack(real_actions, -1).cpu().numpy()
         else:
-            real_actions = torch.cat([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()
+            real_actions = torch.stack([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()
 
         # Single environment step
         obs, reward, done, truncated, _ = env.step(real_actions.reshape(env.action_space.shape))

sheeprl/algos/dreamer_v3/agent.py

@@ -913,7 +913,7 @@ def forward(
                             if sampled_action == 15:  # Craft action
                                 logits[t, b][torch.logical_not(mask["mask_craft_smelt"][t, b])] = -torch.inf
                 elif i == 2:
-                    mask["mask_destroy"][t, b] = mask["mask_destroy"].expand_as(logits)
+                    mask["mask_destroy"] = mask["mask_destroy"].expand_as(logits)
                     mask["mask_equip_place"] = mask["mask_equip_place"].expand_as(logits)
                     for t in range(functional_action.shape[0]):
                         for b in range(functional_action.shape[1]):

sheeprl/algos/dreamer_v3/dreamer_v3.py

@@ -319,7 +319,10 @@ def train(
     critic_grads = None
     if cfg.algo.critic.clip_gradients is not None and cfg.algo.critic.clip_gradients > 0:
         critic_grads = fabric.clip_gradients(
-            module=critic, optimizer=critic_optimizer, max_norm=cfg.algo.critic.clip_gradients, error_if_nonfinite=False
+            module=critic,
+            optimizer=critic_optimizer,
+            max_norm=cfg.algo.critic.clip_gradients,
+            error_if_nonfinite=False,
         )
     critic_optimizer.step()
 
@@ -573,10 +576,10 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
                     real_actions = actions = player.get_actions(torch_obs, mask=mask)
                     actions = torch.cat(actions, -1).cpu().numpy()
                     if is_continuous:
-                        real_actions = torch.cat(real_actions, dim=-1).cpu().numpy()
+                        real_actions = torch.stack(real_actions, dim=-1).cpu().numpy()
                     else:
                         real_actions = (
-                            torch.cat([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()
+                            torch.stack([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()
                         )
 
                 step_data["actions"] = actions.reshape((1, cfg.env.num_envs, -1))

sheeprl/algos/dreamer_v3/utils.py

@@ -125,9 +125,9 @@ def test(
             torch_obs, greedy, {k: v for k, v in torch_obs.items() if k.startswith("mask")}
         )
         if player.actor.is_continuous:
-            real_actions = torch.cat(real_actions, -1).cpu().numpy()
+            real_actions = torch.stack(real_actions, -1).cpu().numpy()
         else:
-            real_actions = torch.cat([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()
+            real_actions = torch.stack([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()
 
         # Single environment step
         obs, reward, done, truncated, _ = env.step(real_actions.reshape(env.action_space.shape))

sheeprl/algos/p2e_dv1/p2e_dv1_exploration.py

@@ -605,7 +605,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
                     real_actions = actions = player.get_exploration_actions(torch_obs, mask=mask, step=policy_step)
                     actions = torch.cat(actions, -1).view(cfg.env.num_envs, -1).cpu().numpy()
                     if is_continuous:
-                        real_actions = torch.cat(real_actions, -1).cpu().numpy()
+                        real_actions = torch.stack(real_actions, -1).cpu().numpy()
                     else:
                         real_actions = (
                             torch.cat([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()

sheeprl/algos/p2e_dv1/p2e_dv1_finetuning.py

@@ -260,10 +260,10 @@ def main(fabric: Fabric, cfg: Dict[str, Any], exploration_cfg: Dict[str, Any]):
                 real_actions = actions = player.get_exploration_actions(torch_obs, mask=mask, step=policy_step)
                 actions = torch.cat(actions, -1).view(cfg.env.num_envs, -1).cpu().numpy()
                 if is_continuous:
-                    real_actions = torch.cat(real_actions, -1).cpu().numpy()
+                    real_actions = torch.stack(real_actions, -1).cpu().numpy()
                 else:
                     real_actions = (
-                        torch.cat([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()
+                        torch.stack([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()
                     )
                 next_obs, rewards, terminated, truncated, infos = envs.step(
                     real_actions.reshape(envs.action_space.shape)

sheeprl/algos/p2e_dv2/p2e_dv2_exploration.py

@@ -128,7 +128,11 @@ def train(
 
     for i in range(0, sequence_length):
         recurrent_state, posterior, prior, posterior_logits, prior_logits = world_model.rssm.dynamic(
-            posterior, recurrent_state, data["actions"][i : i + 1], embedded_obs[i : i + 1], data["is_first"][i : i + 1]
+            posterior,
+            recurrent_state,
+            data["actions"][i : i + 1],
+            embedded_obs[i : i + 1],
+            data["is_first"][i : i + 1],
         )
         recurrent_states[i] = recurrent_state
         priors[i] = prior
@@ -742,10 +746,10 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
                     real_actions = actions = player.get_actions(torch_obs, mask=mask)
                     actions = torch.cat(actions, -1).view(cfg.env.num_envs, -1).cpu().numpy()
                     if is_continuous:
-                        real_actions = torch.cat(real_actions, -1).cpu().numpy()
+                        real_actions = torch.stack(real_actions, -1).cpu().numpy()
                     else:
                         real_actions = (
-                            torch.cat([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()
+                            torch.stack([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()
                         )
 
                 step_data["is_first"] = copy.deepcopy(np.logical_or(step_data["terminated"], step_data["truncated"]))

sheeprl/algos/p2e_dv2/p2e_dv2_finetuning.py

@@ -280,10 +280,10 @@ def main(fabric: Fabric, cfg: Dict[str, Any], exploration_cfg: Dict[str, Any]):
                 real_actions = actions = player.get_actions(torch_obs, mask=mask)
                 actions = torch.cat(actions, -1).view(cfg.env.num_envs, -1).cpu().numpy()
                 if is_continuous:
-                    real_actions = torch.cat(real_actions, -1).cpu().numpy()
+                    real_actions = torch.stack(real_actions, -1).cpu().numpy()
                 else:
                     real_actions = (
-                        torch.cat([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()
+                        torch.stack([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()
                     )
 
                 step_data["is_first"] = copy.deepcopy(np.logical_or(step_data["terminated"], step_data["truncated"]))

sheeprl/algos/p2e_dv3/p2e_dv3_exploration.py

@@ -814,10 +814,10 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
                     real_actions = actions = player.get_actions(torch_obs, mask=mask)
                     actions = torch.cat(actions, -1).cpu().numpy()
                     if is_continuous:
-                        real_actions = torch.cat(real_actions, dim=-1).cpu().numpy()
+                        real_actions = torch.stack(real_actions, dim=-1).cpu().numpy()
                     else:
                         real_actions = (
-                            torch.cat([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()
+                            torch.stack([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()
                         )
 
                 step_data["actions"] = actions.reshape((1, cfg.env.num_envs, -1))

sheeprl/algos/p2e_dv3/p2e_dv3_finetuning.py

@@ -262,10 +262,10 @@ def main(fabric: Fabric, cfg: Dict[str, Any], exploration_cfg: Dict[str, Any]):
                 real_actions = actions = player.get_actions(torch_obs, mask=mask)
                 actions = torch.cat(actions, -1).cpu().numpy()
                 if is_continuous:
-                    real_actions = torch.cat(real_actions, dim=-1).cpu().numpy()
+                    real_actions = torch.stack(real_actions, dim=-1).cpu().numpy()
                 else:
                     real_actions = (
-                        torch.cat([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()
+                        torch.stack([real_act.argmax(dim=-1) for real_act in real_actions], dim=-1).cpu().numpy()
                     )
 
                 step_data["actions"] = actions.reshape((1, cfg.env.num_envs, -1))

sheeprl/algos/ppo/ppo.py

@@ -276,9 +276,9 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
                     )
                     actions, logprobs, values = player(torch_obs)
                     if is_continuous:
-                        real_actions = torch.cat(actions, -1).cpu().numpy()
+                        real_actions = torch.stack(actions, -1).cpu().numpy()
                     else:
-                        real_actions = torch.cat([act.argmax(dim=-1) for act in actions], dim=-1).cpu().numpy()
+                        real_actions = torch.stack([act.argmax(dim=-1) for act in actions], dim=-1).cpu().numpy()
                     actions = torch.cat(actions, -1).cpu().numpy()
 
                     # Single environment step

sheeprl/algos/ppo/ppo_decoupled.py

@@ -206,9 +206,9 @@ def player(
                 torch_obs = prepare_obs(fabric, next_obs, cnn_keys=cfg.algo.cnn_keys.encoder, num_envs=cfg.env.num_envs)
                 actions, logprobs, values = agent(torch_obs)
                 if is_continuous:
-                    real_actions = torch.cat(actions, -1).cpu().numpy()
+                    real_actions = torch.stack(actions, -1).cpu().numpy()
                 else:
-                    real_actions = torch.cat([act.argmax(dim=-1) for act in actions], dim=-1).cpu().numpy()
+                    real_actions = torch.stack([act.argmax(dim=-1) for act in actions], dim=-1).cpu().numpy()
                 actions = torch.cat(actions, -1).cpu().numpy()
 
                 # Single environment step

sheeprl/algos/ppo_recurrent/ppo_recurrent.py

@@ -299,9 +299,9 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
                         torch_obs, prev_actions=torch_prev_actions, prev_states=prev_states
                     )
                     if is_continuous:
-                        real_actions = torch.cat(actions, -1).cpu().numpy()
+                        real_actions = torch.stack(actions, -1).cpu().numpy()
                     else:
-                        real_actions = torch.cat([act.argmax(dim=-1) for act in actions], dim=-1).cpu().numpy()
+                        real_actions = torch.stack([act.argmax(dim=-1) for act in actions], dim=-1).cpu().numpy()
                     torch_actions = torch.cat(actions, dim=-1)
                     actions = torch_actions.cpu().numpy()
 

sheeprl/algos/ppo_recurrent/utils.py

@@ -57,10 +57,10 @@ def test(agent: "RecurrentPPOPlayer", fabric: Fabric, cfg: Dict[str, Any], log_d
         # Act greedly through the environment
         actions, state = agent.get_actions(torch_obs, actions, state, greedy=True)
         if agent.actor.is_continuous:
-            real_actions = torch.cat(actions, -1)
+            real_actions = torch.stack(actions, -1)
             actions = torch.cat(actions, dim=-1).view(1, 1, -1)
         else:
-            real_actions = torch.cat([act.argmax(dim=-1) for act in actions], dim=-1)
+            real_actions = torch.stack([act.argmax(dim=-1) for act in actions], dim=-1)
             actions = torch.cat([act for act in actions], dim=-1).view(1, 1, -1)
 
         # Single environment step

sheeprl/algos/sac/sac_decoupled.py

@@ -189,7 +189,7 @@ def player(
                 torch_obs = prepare_obs(fabric, obs, num_envs=cfg.env.num_envs)
                 actions = actor(torch_obs)
                 actions = actions.cpu().numpy()
-            next_obs, rewards, terminated, truncated, infos = envs.step(actions)
+            next_obs, rewards, terminated, truncated, infos = envs.step(actions.reshape(envs.action_space.shape))
             rewards = rewards.reshape(cfg.env.num_envs, -1)
 
         if cfg.metric.log_level > 0 and "final_info" in infos:

sheeprl/configs/exp/dreamer_v3_minedojo.yaml

@@ -0,0 +1,57 @@
+# @package _global_
+
+defaults:
+  - dreamer_v3
+  - override /algo: dreamer_v3_XS
+  - override /env: minedojo
+  - _self_
+
+# Experiment
+seed: 5
+total_steps: 50000000
+
+# Environment
+env:
+  num_envs: 2
+  id: harvest_milk
+  reward_as_observation: True
+
+# Checkpoint
+checkpoint:
+  every: 100000
+
+# Buffer
+buffer:
+  checkpoint: True
+
+# Algorithm
+algo:
+  replay_ratio: 0.015625
+  learning_starts: 65536
+  actor:
+    cls: sheeprl.algos.dreamer_v3.agent.MinedojoActor
+  cnn_keys:
+    encoder:
+      - rgb
+  mlp_keys:
+    encoder:
+      - equipment
+      - inventory
+      - inventory_delta
+      - inventory_max
+      - life_stats
+      - mask_action_type
+      - mask_craft_smelt
+      - mask_destroy
+      - mask_equip_place
+      - reward
+    decoder:
+      - equipment
+      - inventory
+      - inventory_delta
+      - inventory_max
+      - life_stats
+      - mask_action_type
+      - mask_craft_smelt
+      - mask_destroy
+      - mask_equip_place