benign_random. Add apm to args

stratosphereips · Apr 25, 2024 · b83636e · b83636e
1 parent 0837c12
commit b83636e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/agents/benign_random/benign_random_agent.py b/agents/benign_random/benign_random_agent.py
@@ -79,12 +79,13 @@ def select_action(self, observation:Observation)->Action:
     parser.add_argument("--port", help="Port where the game server is", default=9000, type=int, action='store', required=False)
     parser.add_argument("--episodes", help="Sets number of testing episodes", default=10, type=int)
     parser.add_argument("--logdir", help="Folder to store logs", default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "logs"))
+    parser.add_argument("--apm", help="Actions per minute", default=10, type=int, required=False)
     args = parser.parse_args()
 
     if not os.path.exists(args.logdir):
         os.makedirs(args.logdir)
     logging.basicConfig(filename=os.path.join(args.logdir, "benign_random_agent.log"), filemode='w', format='%(asctime)s %(name)s %(levelname)s %(message)s', datefmt='%H:%M:%S',level=logging.DEBUG)
 
     # Create agent
-    agent = RandomBenignAgent(args.host, args.port,"Human",allowed_actions=[ActionType.FindData, ActionType.ExfiltrateData, ActionType.FindServices], apm_limit=10)
+    agent = RandomBenignAgent(args.host, args.port,"Human",allowed_actions=[ActionType.FindData, ActionType.ExfiltrateData, ActionType.FindServices], apm_limit=args.apm)
     agent.play_game(args.episodes)