moved plot outside and copied the get test rollout code which worked

Author: miTTimmiTTim

ddopai/_modidx.py

@@ -1745,7 +1745,9 @@
                                                                            '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.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/pricing_env/pricing_env.py'),
+                                                                           'ddopai.meta_learning.environments.pricing_env.pricing_env.plot_actions_reward': ( '50_meta_learning/53_environments/01_pricing_env/pricing_env.html#plot_actions_reward',
+                                                                                                                                                              '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__ = ['device', 'PricingEnv']
+__all__ = ['device', 'PricingEnv', 'plot_actions_reward']
 
 # %% ../../../../nbs/50_meta_learning/53_environments/01_pricing_env/10_pricing_env.ipynb 1
 import gym
@@ -261,153 +261,141 @@ def visualise_behaviour(env,
         """
 
         num_episodes = args.max_rollouts_per_task
-        unwrapped_env = env.venv.unwrapped.envs[0]
 
-        episode_all_obs = [[] for _ in range(num_episodes)]
+        # --- initialise things we want to keep track of ---
+
         episode_prev_obs = [[] for _ in range(num_episodes)]
         episode_next_obs = [[] for _ in range(num_episodes)]
         episode_actions = [[] for _ in range(num_episodes)]
         episode_rewards = [[] for _ in range(num_episodes)]
+
         episode_returns = []
         episode_lengths = []
 
-        if args.pass_belief_to_policy and (encoder is None):
-            episode_beliefs = [[] for _ in range(num_episodes)]
-        else:
-            episode_beliefs = None
-            
         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
+            curr_latent_sample = curr_latent_mean = curr_latent_logvar = None
+            episode_latent_means = episode_latent_logvars = None
+
+        # --- roll out policy ---
 
+        # (re)set environment
         env.reset_task()
-        [state, belief, task] = utl.reset_env(env, args)
-        start_obs = state.clone()
+        state, belief, task = utl.reset_env(env, args)
+        state = state.reshape((1, -1)).to(device)
+        task = task.view(-1) if task is not None else None
 
         for episode_idx in range(num_episodes):
 
             curr_rollout_rew = []
 
-            
-
             if encoder is not None:
-                
-                if episode_idx == 0 and encoder is not None:
+                if episode_idx == 0:
                     # reset to prior
                     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)
-                    
                 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_all_obs[episode_idx].append(start_obs.clone())
-            if args.pass_belief_to_policy and (encoder is None):
-                episode_beliefs[episode_idx].append(belief)
-                
-            for step_idx in range(env._max_episode_steps):
-                
-                if step_idx == 1:
-                    prev_obs = start_obs.clone()
-                else:
-                    prev_obs = state.clone()
-                    
-                episode_prev_obs[episode_idx].append(prev_obs)
-                
-                                # act
-                _, action, _ = utl.select_action(args=args,
-                                                 policy=policy,
-                                                 state=state.view(-1),
-                                                 belief=belief,
-                                                 task=task,
-                                                 deterministic=True,
-                                                 latent_sample=curr_latent_sample.view(-1) if (curr_latent_sample is not None) else None,
-                                                 latent_mean=curr_latent_mean.view(-1) if (curr_latent_mean is not None) else None,
-                                                 latent_logvar=curr_latent_logvar.view(-1) if (curr_latent_logvar is not None) else None,
-                                                 )
-                
+            for step_idx in range(1, env._max_episode_steps + 1):
+
+                episode_prev_obs[episode_idx].append(state.clone())
+                prev_state = state.clone()
+
+                latent = utl.get_latent_for_policy(args,
+                                                latent_sample=curr_latent_sample,
+                                                latent_mean=curr_latent_mean,
+                                                latent_logvar=curr_latent_logvar)
+                _, action, _ = policy.act(state=state.view(-1), latent=latent, belief=belief, task=task, deterministic=True)
+                action = action.reshape((1, *action.shape))
+
                 # observe reward and next obs
-                [state, belief, task], (rew_raw, rew_normalised), done, infos = utl.env_step(env, action, args)
+                (state, belief, task), (rew_raw, rew_normalised), done, infos = utl.env_step(env, action, args)
+                state = state.reshape((1, -1)).to(device)
+                task = task.view(-1) if task is not None else None
 
                 if encoder is not None:
                     # update task embedding
                     curr_latent_sample, curr_latent_mean, curr_latent_logvar, hidden_state = encoder(
                         action.float().to(device),
                         state,
                         rew_raw.reshape((1, 1)).float().to(device),
-                        prev_obs,
+                        prev_state,
                         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_all_obs[episode_idx].append(state.clone())
                 episode_next_obs[episode_idx].append(state.clone())
                 episode_rewards[episode_idx].append(rew_raw.clone())
                 episode_actions[episode_idx].append(action.clone())
 
-                curr_rollout_rew.append(rew_raw.clone())
-                
-
-                if args.pass_belief_to_policy and (encoder is None):
-                    episode_beliefs[episode_idx].append(belief)
-
-                if infos[0]['done_mdp'] and not done:
-                    start_obs = infos[0]['start_state']
-                    start_obs = torch.from_numpy(start_obs).float().reshape((1, -1)).to(device)
+                if infos[0]['done_mdp']:
                     break
 
             episode_returns.append(sum(curr_rollout_rew))
             episode_lengths.append(step_idx)
 
-
         # clean up
-
         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.cat(e) for e in episode_actions]
-        episode_rewards = [torch.cat(e) for e in episode_rewards]
-
-
-        # Plot price and reward trajectories
-        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()
-
+        episode_rewards = [torch.cat(r) for r in episode_rewards]
+
+        plot_actions_reward(
+            episode_actions=episode_actions,
+            episode_rewards=episode_rewards,
+            episode_lengths=episode_lengths,
+            image_folder=image_folder,
+            iter_idx=iter_idx
+        )
         return episode_latent_means, episode_latent_logvars, \
-            episode_prev_obs, episode_next_obs, episode_actions, episode_rewards, \
-            episode_returns
+           episode_prev_obs, episode_next_obs, episode_actions, episode_rewards, \
+           episode_returns
+
+
+        
 
 
 
 
 
+
+# %% ../../../../nbs/50_meta_learning/53_environments/01_pricing_env/10_pricing_env.ipynb 3
+def plot_actions_reward(
+    episode_actions: List[torch.Tensor],
+    episode_rewards: List[torch.Tensor],
+    episode_lengths: List[int],
+    image_folder: Optional[str] = None,
+    iter_idx: int = 0
+):
+    """
+    Plot actions and rewards for each episode.
+    """
+    plt.figure(figsize=(12, 6))
+    for i, (actions, rewards, length) in enumerate(zip(episode_actions, episode_rewards, episode_lengths)):
+        plt.plot(range(length), actions.cpu().numpy(), label=f'Episode {i+1} Actions')
+        plt.plot(range(length), rewards.cpu().numpy(), label=f'Episode {i+1} Rewards', linestyle='--')
+
+    plt.xlabel('Time Step')
+    plt.ylabel('Value')
+    plt.title('Actions and Rewards per Episode')
+    plt.legend()
+    
+    if image_folder:
+        plt.savefig(f"{image_folder}/actions_rewards_iter_{iter_idx}.png")
+    else:
+        plt.show()  
+