Updating the trainer UI

app.py

@@ -9,30 +9,6 @@ import transformers
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-#
-# model = "tiiuae/falcon-7b-instruct"
-#
-# tokenizer = AutoTokenizer.from_pretrained(model)
-# pipeline = transformers.pipeline(
-#     "text-generation",
-#     model=model,
-#     tokenizer=tokenizer,
-#     torch_dtype=torch.bfloat16,
-#     trust_remote_code=True,
-#     device_map="auto",
-# )
-# sequences = pipeline(
-#    "Girafatron is obsessed with giraffes, the most glorious animal on the face of this Earth. Giraftron believes all other animals are irrelevant when compared to the glorious majesty of the giraffe.\nDaniel: Hello, Girafatron!\nGirafatron:",
-#     max_length=200,
-#     do_sample=True,
-#     top_k=10,
-#     num_return_sequences=1,
-#     eos_token_id=tokenizer.eos_token_id,
-# )
-# st.title("Beyond the Anti-Jam: Integration of DRL with LLM")
-# for seq in sequences:
-#     st.write(f"Result: {seq['generated_text']}")
-
 
 def perform_training(jammer_type, channel_switching_cost):
     agent = train(jammer_type, channel_switching_cost)
@@ -68,8 +44,31 @@ st.sidebar.write(f"Channel Switching Cost: {channel_switching_cost}")
 start_button = st.sidebar.button('Start')
 
 if start_button:
-    agent = perform_training(jammer_type, channel_switching_cost)
+    agent, rewards = perform_training(jammer_type, channel_switching_cost)
     st.subheader("Generating Insights of the DRL-Training")
     # text = pipeline("Discuss this topic: Integrating LLMs to DRL-based anti-jamming.")
     # st.write(text)
     test(agent, jammer_type, channel_switching_cost)
+
+# model = "tiiuae/falcon-7b-instruct"
+#
+# tokenizer = AutoTokenizer.from_pretrained(model)
+# pipeline = transformers.pipeline(
+#     "text-generation",
+#     model=model,
+#     tokenizer=tokenizer,
+#     torch_dtype=torch.bfloat16,
+#     trust_remote_code=True,
+#     device_map="auto",
+# )
+# sequences = pipeline(
+#    "Girafatron is obsessed with giraffes, the most glorious animal on the face of this Earth. Giraftron believes all other animals are irrelevant when compared to the glorious majesty of the giraffe.\nDaniel: Hello, Girafatron!\nGirafatron:",
+#     max_length=200,
+#     do_sample=True,
+#     top_k=10,
+#     num_return_sequences=1,
+#     eos_token_id=tokenizer.eos_token_id,
+# )
+# st.title("Beyond the Anti-Jam: Integration of DRL with LLM")
+# for seq in sequences:
+#     st.write(f"Result: {seq['generated_text']}")

trainer.py

@@ -79,10 +79,23 @@ def train(jammer_type, channel_switching_cost):
     ax.set_title(f'Training Rewards - {jammer_type}, CSC: {channel_switching_cost}')
     ax.legend()
 
-    # Display the Streamlit figure using streamlit.pyplot
-    st.set_option('deprecation.showPyplotGlobalUse', False)
-    st.subheader("Training Graph")
-    st.pyplot(fig)
+    # Use Streamlit layout to create two side-by-side containers
+    with st.beta_container():
+        col1, col2 = st.beta_columns(2)
+
+        with col1:
+            st.subheader("Training Graph")
+            st.set_option('deprecation.showPyplotGlobalUse', False)
+            st.pyplot(fig)
+
+        with col2:
+            st.subheader("Graph Explanation")
+            st.write("""
+               The training graph shows the rewards received by the agent in each episode of the training process.
+               The blue line represents the actual reward values, while the black line represents a rolling average.
+               The red horizontal line indicates the threshold for considering the task solved.
+               The green line represents the epsilon (exploration rate) values for the agent, indicating how often it takes random actions.
+               """)
 
     # Save the figure
     # plot_name = f'./data/train_rewards_{jammer_type}_csc_{channel_switching_cost}.png'
@@ -98,4 +111,4 @@ def train(jammer_type, channel_switching_cost):
     # # Save the agent as a SavedAgent.
     # agentName = f'./data/DDQNAgent_{jammer_type}_csc_{channel_switching_cost}'
     # DDQN_agent.save_model(agentName)
-    return DDQN_agent
+    return DDQN_agent, rewards