Fixed horizon call in visualise env

Author: miTTimmiTTim

ddopai/_modidx.py

@@ -1744,8 +1744,8 @@
                                                                                                                                                           'ddopai/meta_learning/environments/pricing_env/pricing_env.py'),
                                                                            'ddopai.meta_learning.environments.pricing_env.pricing_env.PricingEnv.step': ( '50_meta_learning/53_environments/01_pricing_env/pricing_env.html#pricingenv.step',
                                                                                                                                                           'ddopai/meta_learning/environments/pricing_env/pricing_env.py'),
-                                                                           'ddopai.meta_learning.environments.pricing_env.pricing_env.visualise_behaviour': ( '50_meta_learning/53_environments/01_pricing_env/pricing_env.html#visualise_behaviour',
-                                                                                                                                                              'ddopai/meta_learning/environments/pricing_env/pricing_env.py')},
+                                                                           'ddopai.meta_learning.environments.pricing_env.pricing_env.PricingEnv.visualise_behaviour': ( '50_meta_learning/53_environments/01_pricing_env/pricing_env.html#pricingenv.visualise_behaviour',
+                                                                                                                                                                         'ddopai/meta_learning/environments/pricing_env/pricing_env.py')},
             'ddopai.meta_learning.environments.wrappers': { 'ddopai.meta_learning.environments.wrappers.PrevActRewWrapper': ( '50_meta_learning/53_environments/wrappers.html#prevactrewwrapper',
                                                                                                                               'ddopai/meta_learning/environments/wrappers.py'),
                                                             'ddopai.meta_learning.environments.wrappers.PrevActRewWrapper.__init__': ( '50_meta_learning/53_environments/wrappers.html#prevactrewwrapper.__init__',

ddopai/meta_learning/environments/pricing_env/pricing_env.py

@@ -1,7 +1,7 @@
 # AUTOGENERATED! DO NOT EDIT! File to edit: ../../../../nbs/50_meta_learning/53_environments/01_pricing_env/10_pricing_env.ipynb.
 
 # %% auto 0
-__all__ = ['PricingEnv', 'visualise_behaviour']
+__all__ = ['device', 'PricingEnv']
 
 # %% ../../../../nbs/50_meta_learning/53_environments/01_pricing_env/10_pricing_env.ipynb 1
 import gym
@@ -11,6 +11,7 @@
 import matplotlib.pyplot as plt
 import torch
 from  ...utils import helpers as utl
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 # %% ../../../../nbs/50_meta_learning/53_environments/01_pricing_env/10_pricing_env.ipynb 2
 class PricingEnv(gym.Env):
@@ -245,135 +246,129 @@ def _demand(self, price: float, noise: float) -> float:
         return max(0.0, mean + noise)
 
     # ---------- visualisation stub -------------------------------------------
-@staticmethod
-def visualise_behaviour(env,
-                        args,
-                        policy,
-                        iter_idx,
-                        encoder=None,
-                        image_folder=None,
-                        return_pos=False,
-                        **kwargs):
-
-    num_episodes = args.max_rollouts_per_task
-
-    episode_prev_obs = [[] for _ in range(num_episodes)]
-    episode_next_obs = [[] for _ in range(num_episodes)]
-    episode_actions  = [[] for _ in range(num_episodes)]  # price = action
-    episode_rewards  = [[] for _ in range(num_episodes)]
-    episode_returns  = []
-
-    if encoder is not None:
-        episode_latent_samples = [[] for _ in range(num_episodes)]
-        episode_latent_means   = [[] for _ in range(num_episodes)]
-        episode_latent_logvars = [[] for _ in range(num_episodes)]
-    else:
-        episode_latent_samples = episode_latent_means = episode_latent_logvars = None
-
-    env.reset_task()
-    state, belief, task = utl.reset_env(env, args)
-    task = task.view(-1) if task is not None else None
-
-    hidden_state = torch.zeros((1, args.hidden_size)).to(args.device) if hasattr(args, 'hidden_size') else None
-
-    for episode_idx in range(num_episodes):
-        curr_rollout_rew = []
-
-        if episode_idx == 0:
-            if encoder is not None:
-                curr_latent_sample, curr_latent_mean, curr_latent_logvar, hidden_state = encoder.prior(1)
-                curr_latent_sample = curr_latent_sample[0].to(args.device)
-                curr_latent_mean   = curr_latent_mean[0].to(args.device)
-                curr_latent_logvar = curr_latent_logvar[0].to(args.device)
-            else:
-                curr_latent_sample = curr_latent_mean = curr_latent_logvar = None
+    @staticmethod
+    def visualise_behaviour(env,
+                            args,
+                            policy,
+                            iter_idx,
+                            encoder=None,
+                            image_folder=None,
+                            **kwargs):
+
+        num_episodes = args.max_rollouts_per_task
+
+        episode_prev_obs = [[] for _ in range(num_episodes)]
+        episode_next_obs = [[] for _ in range(num_episodes)]
+        episode_actions  = [[] for _ in range(num_episodes)]  # price = action
+        episode_rewards  = [[] for _ in range(num_episodes)]
+        episode_returns  = []
 
         if encoder is not None:
-            episode_latent_samples[episode_idx].append(curr_latent_sample[0].clone())
-            episode_latent_means[episode_idx].append(curr_latent_mean[0].clone())
-            episode_latent_logvars[episode_idx].append(curr_latent_logvar[0].clone())
+            episode_latent_samples = [[] for _ in range(num_episodes)]
+            episode_latent_means   = [[] for _ in range(num_episodes)]
+            episode_latent_logvars = [[] for _ in range(num_episodes)]
+        else:
+            episode_latent_samples = episode_latent_means = episode_latent_logvars = None
+
+        env.reset_task()
+        state, belief, task = utl.reset_env(env, args)
+        task = task.view(-1) if task is not None else None
+
+        hidden_state = torch.zeros((1, args.hidden_size)).to(device) if hasattr(args, 'hidden_size') else None
+
+        for episode_idx in range(num_episodes):
+            curr_rollout_rew = []
+
+            if episode_idx == 0:
+                if encoder is not None:
+                    curr_latent_sample, curr_latent_mean, curr_latent_logvar, hidden_state = encoder.prior(1)
+                    curr_latent_sample = curr_latent_sample[0].to(device)
+                    curr_latent_mean   = curr_latent_mean[0].to(device)
+                    curr_latent_logvar = curr_latent_logvar[0].to(device)
+                else:
+                    curr_latent_sample = curr_latent_mean = curr_latent_logvar = None
 
-        obs = env.reset()
+            if encoder is not None:
+                episode_latent_samples[episode_idx].append(curr_latent_sample[0].clone())
+                episode_latent_means[episode_idx].append(curr_latent_mean[0].clone())
+                episode_latent_logvars[episode_idx].append(curr_latent_logvar[0].clone())
 
-        for step_idx in range(1, env.horizon + 1):
-            prev_obs = torch.tensor(obs, dtype=torch.float32).to(args.device).unsqueeze(0)
-            episode_prev_obs[episode_idx].append(prev_obs.clone())
+            obs = env.reset()
 
-            latent = utl.get_latent_for_policy(args,
-                                               latent_sample=curr_latent_sample,
-                                               latent_mean=curr_latent_mean,
-                                               latent_logvar=curr_latent_logvar)
+            for step_idx in range(1, env._max_episode_steps + 1):
+                prev_obs = torch.tensor(obs, dtype=torch.float32).to(device).unsqueeze(0)
+                episode_prev_obs[episode_idx].append(prev_obs.clone())
 
-            _, action, _ = policy.act(prev_obs, latent, belief=None, task=task, deterministic=True)
+                latent = utl.get_latent_for_policy(args,
+                                                latent_sample=curr_latent_sample,
+                                                latent_mean=curr_latent_mean,
+                                                latent_logvar=curr_latent_logvar)
 
-            obs, reward, done, info = env.step(action.cpu().numpy())
-            obs = torch.tensor(obs, dtype=torch.float32).to(args.device).unsqueeze(0)
+                _, action, _ = policy.act(prev_obs, latent, belief=None, task=task, deterministic=True)
 
-            episode_next_obs[episode_idx].append(obs.clone())
-            episode_actions[episode_idx].append(action.clone())
-            episode_rewards[episode_idx].append(torch.tensor([reward], dtype=torch.float32).to(args.device))
-            curr_rollout_rew.append(reward)
+                obs, reward, done, info = env.step(action.cpu().numpy())
+                obs = torch.tensor(obs, dtype=torch.float32).to(device).unsqueeze(0)
 
-            if encoder is not None:
-                curr_latent_sample, curr_latent_mean, curr_latent_logvar, hidden_state = encoder(
-                    action.reshape(1, -1).float().to(args.device),
-                    obs,
-                    torch.tensor([reward], dtype=torch.float32, device=args.device).reshape(1, -1),
-                    prev_obs,
-                    hidden_state,
-                    return_prior=False,
-                )
-                episode_latent_samples[episode_idx].append(curr_latent_sample[0].clone())
-                episode_latent_means[episode_idx].append(curr_latent_mean[0].clone())
-                episode_latent_logvars[episode_idx].append(curr_latent_logvar[0].clone())
+                episode_next_obs[episode_idx].append(obs.clone())
+                episode_actions[episode_idx].append(action.clone())
+                episode_rewards[episode_idx].append(torch.tensor([reward], dtype=torch.float32).to(device))
+                curr_rollout_rew.append(reward)
 
-            if done:
-                break
-
-        episode_returns.append(sum(curr_rollout_rew))
-
-    # Convert to tensor batches
-    if encoder is not None:
-        episode_latent_means = [torch.stack(e) for e in episode_latent_means]
-        episode_latent_logvars = [torch.stack(e) for e in episode_latent_logvars]
-
-    episode_prev_obs = [torch.cat(e) for e in episode_prev_obs]
-    episode_next_obs = [torch.cat(e) for e in episode_next_obs]
-    episode_actions  = [torch.stack(e) for e in episode_actions]
-    episode_rewards  = [torch.cat(e) for e in episode_rewards]
-
-    # ---- Plot: Price (action) and Revenue ----
-    import matplotlib.pyplot as plt
-
-    plt.figure(figsize=(10, 3 * num_episodes))
-    for i in range(num_episodes):
-        plt.subplot(num_episodes, 2, 2 * i + 1)
-        plt.plot(episode_actions[i].cpu().numpy(), label="Price")
-        plt.ylabel("Price")
-        plt.xlabel("Timestep")
-        plt.title(f"Episode {i}: Price")
-
-        plt.subplot(num_episodes, 2, 2 * i + 2)
-        plt.plot(episode_rewards[i].cpu().numpy(), label="Revenue", color='green')
-        plt.ylabel("Revenue")
-        plt.xlabel("Timestep")
-        plt.title(f"Episode {i}: Revenue")
-
-    plt.tight_layout()
-    if image_folder is not None:
-        plt.savefig(f"{image_folder}/{iter_idx}_pricing_behaviour.png")
-        plt.close()
-    else:
-        plt.show()
-
-    if not return_pos:
-        return episode_latent_means, episode_latent_logvars, \
-               episode_prev_obs, episode_next_obs, episode_actions, episode_rewards, \
-               episode_returns
-    else:
+                if encoder is not None:
+                    curr_latent_sample, curr_latent_mean, curr_latent_logvar, hidden_state = encoder(
+                        action.reshape(1, -1).float().to(device),
+                        obs,
+                        torch.tensor([reward], dtype=torch.float32, device=device).reshape(1, -1),
+                        prev_obs,
+                        hidden_state,
+                        return_prior=False,
+                    )
+                    episode_latent_samples[episode_idx].append(curr_latent_sample[0].clone())
+                    episode_latent_means[episode_idx].append(curr_latent_mean[0].clone())
+                    episode_latent_logvars[episode_idx].append(curr_latent_logvar[0].clone())
+
+                if done:
+                    break
+
+            episode_returns.append(sum(curr_rollout_rew))
+
+        # Convert to tensor batches
+        if encoder is not None:
+            episode_latent_means = [torch.stack(e) for e in episode_latent_means]
+            episode_latent_logvars = [torch.stack(e) for e in episode_latent_logvars]
+
+        episode_prev_obs = [torch.cat(e) for e in episode_prev_obs]
+        episode_next_obs = [torch.cat(e) for e in episode_next_obs]
+        episode_actions  = [torch.stack(e) for e in episode_actions]
+        episode_rewards  = [torch.cat(e) for e in episode_rewards]
+
+
+        plt.figure(figsize=(10, 3 * num_episodes))
+        for i in range(num_episodes):
+            plt.subplot(num_episodes, 2, 2 * i + 1)
+            plt.plot(episode_actions[i].cpu().numpy(), label="Price")
+            plt.ylabel("Price")
+            plt.xlabel("Timestep")
+            plt.title(f"Episode {i}: Price")
+
+            plt.subplot(num_episodes, 2, 2 * i + 2)
+            plt.plot(episode_rewards[i].cpu().numpy(), label="Revenue", color='green')
+            plt.ylabel("Revenue")
+            plt.xlabel("Timestep")
+            plt.title(f"Episode {i}: Revenue")
+
+        plt.tight_layout()
+        if image_folder is not None:
+            plt.savefig(f"{image_folder}/{iter_idx}_pricing_behaviour.png")
+            plt.close()
+        else:
+            plt.show()
+
+        
         return episode_latent_means, episode_latent_logvars, \
-               episode_prev_obs, episode_next_obs, episode_actions, episode_rewards, \
-               episode_returns, episode_actions  # actions = price = pos
+                episode_prev_obs, episode_next_obs, episode_actions, episode_rewards, \
+                episode_returns
+