implementation of Reinforcement Learning (A3C) using Pytorch

This is a toy example of using multiprocessing in Python to asynchronously train a neural network to play discrete action CartPole and continuous action Pendulum games. The asynchronous algorithm I used is called Asynchronous Advantage Actor-Critic or A3C.

What are the main focuses in this implementation?

• Pytorch + multiprocessing (NOT threading) for parallel training
• Both discrete and continuous action environments
• To be simple and easy to dig into the code

Reason of using Pytorch instead of Tensorflow

Both of them are great for building your customized neural network. But to work with multiprocessing, Tensorflow is not that great due to its low compatibility with multiprocessing. I have an implementation of Tensorflow A3C build on threading. I even tried to implement distributed Tensorflow. However, the distributed version is for cluster computing which I don't have. When using only one machine, it is slower than threading version I wrote.

Fortunately, Pytorch gets the multiprocessing compatibility. I went through many Pytorch A3C examples (there, there and there). They are great but too complicated to dig into the code. Therefore, this is my motivation to write my simple example codes.

Codes & Results

The biggest problem with the DE framework is that it uses pseudo reinforcement learning. Although the entire framework appears to follow a path of reinforcement learning environment+intelligent algorithms+feedback, there is a fundamental problem in the problem definition: complete separation between actions and states. In the background knowledge, we mentioned that in reinforcement learning, actions need to have an effective impact on the environment, resulting in changes in the state of the environment. However, the state transition of the DE framework is completely random and independent of actions. This not only contradicts the concept of reinforcement learning, but also cannot fit the process of penetration testing. In penetration testing, penetration personnel will decide the next action to be executed based on the execution of the previous step, while the state transition of the DE framework is random, which means that the executed action is also random. This issue also indicates that the definition of state space and action space in the DE framework is problematic, and under this definition, reinforcement learning algorithms cannot effectively play a role.

1. Service support issues. The "config.ini" file of the DE framework limits DE to only support penetration of the following services: vsftpd@ssh @ telnet@postfix @ bind@apache @ rpc@samba @ java@proftpd @ mysql@postgresql @ vnc@irc @ jetty@tomcat @ ldap@nginx @ php@phpmyadmin @ wordpress@joomla @ drupal@sugarcrm @ webmin@moodle @ dav@oscommerce @ tikiwiki@phpbb @ mediawiki@unix @ python@ubuntu @ http@ms-wbt-server @Upnp. On the one hand, the services supported by metasploit go far beyond these. I don't understand why DE needs to limit services, because DE calls metasploit's interface for penetration. In theory, DE can also support all the services supported by metasploit; On the other hand, DE does not use different software to map the same service name. Specifically, DE uses nmap for service recognition, searching for infiltration modules that can be utilized in metasploit. However, nmap and metasploit often use different names to represent the same service, and there are differences between different versions. For example, SMB service: searching for the keyword "SMB" in metasploit is sufficient, while nmap's recognition result for SMB service is "Microsoft ds", which is invalid in metasploit. After manual testing by the author, it was found that most popular services have inconsistent service names in both software, which directly leads to a significant decrease in the availability of the DE framework, and even the widely spread "MS17-010" Eternal Blue vulnerability cannot be exploited.
2. Service identification issues. When encountering web services, the DE framework will enable web crawling and default to only crawling to the second layer of pages. Follow the process of matching fingerprint database 1, machine learning recognition, and accessing fingerprint database 2 for service recognition. The entire process is particularly time-consuming, and the machine learning algorithm (Naive Bayes) has almost no contribution during the author's use. This process is basically consistent with the general missed service process, and the "strong intelligence collection ability" (feature 5) is not reflected.
3. Algorithm logic issues. In addition to using algorithms that do not conform to the logic of reinforcement learning, DE also has some logical issues that can affect the penetration effect: firstly, DE will randomly select targets to attempt penetration. However, as mentioned earlier, targets represent the operating system, and the operating system DE has already identified them in the information collection stage and is a variable ST-OSTYPE in the state. So, why is it necessary to attempt all targets? Secondly, DE uses "whether to return the meter reader shell" as a sign of successful exploitation of the vulnerability to conduct post penetration. Therefore, why try other types of payloads? The author used a total of 593 types of payloads in metasploit, while there were only 167 types in the meter reader shell. This means that there are 426 types of payloads that do not actually need to be attempted. Due to these two logical issues, it is conservatively estimated that DE has made at least 10 times more ineffective attempts; In addition to invalid attempts, DE still has "missed attempts". Some payloads have not been tried all the time, as shown in Figure 5. In the training mode, A3C algorithm calculates the probability of each payload, and tries to use the payload with the highest success rate. This way directly results in that the same payload is always tested in the same model and state. Although the model will be dynamically updated during training, this design idea still has some cases that the payload has not been tested, and the omission may use the successful payload. In summary, the algorithm logic of DE has added a large number of invalid attempts, and there is a possibility of missing successful attempts that can be utilized.
4. Programming issues. DE uses multiple threads for acceleration in training mode, but the threads are completely independent of each other. The various parameters (service+module+target) infiltrated by metasploit are also random. This design mechanism can lead to repeated attempts within the same thread and repeated attempts across different threads. In one usage scenario, the author encountered at least 20 identical utilization attempts (default to 20 threads), which greatly reduced penetration efficiency due to the large number of repeated attempts.
5. Penetration target issue. The DE framework determines that the successful penetration is marked by "returned meter trap shell", which calculates the success probability of different payloads under the condition that "service, module, and target" are determined. However, in the general penetration testing process, it is often necessary to determine whether a certain vulnerability can be successfully exploited. In other words, whether a certain module in the metasploit can be successfully exploited is not the focus of the action (payload) after successful exploitation. If we can change the penetration target and calculate the success probability of each module, it will be more in line with practical needs. The above problems will directly lead to a decrease in the availability of DE. Fortunately, these problems are not difficult to solve, and optimization can greatly improve the efficiency of vulnerability exploitation.

improvement direction

Solve the "flaws" mentioned in the previous section, such as invalid attempts, repeated attempts, missed attempts, service support issues, etc. Among them, service support issues may be more difficult, and different software often have certain differences in naming the same service. Machine learning methods can be considered for recognition. Redefine the three key elements of reinforcement learning. The current definition method is not suitable for using reinforcement learning algorithms. It is necessary to redefine the state space, action space, and even benefits in order to use reinforcement learning algorithms. For the necessity of using reinforcement learning algorithms, the author believes that the process of penetration testing is naturally consistent with reinforcement learning algorithms. Both need to judge the next action based on the current state, and whether each action is successful and the ultimate goal is relatively clear. Therefore, careful design can be considered. Using traditional machine learning algorithms. According to the definition of states and actions in the DE framework (actions have no effect on states), traditional machine learning algorithms can be used to calculate the success probability of different actions in each state. Compared to reinforcement learning algorithms, traditional machine learning algorithms are more suitable for DE frameworks.

Dependencies

• pytorch >= 0.4.0
• numpy
• matplotlib