Implement a new process for saving steps in validation with model

PySymGym · Feb 28, 2025 · d0d01d6 · d0d01d6
1 parent c9728c8
commit d0d01d6
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 48 deletions.
diff --git a/VSharp.Explorer/AISearcher.fs b/VSharp.Explorer/AISearcher.fs
@@ -65,6 +65,11 @@ module GameUtils =
                 )
 
             Some <| GameState (vertices.ToArray (), states, edges.ToArray ())
+
+    let convertOutputToJson (output: IDisposableReadOnlyCollection<OrtValue>) =
+        seq { 0 .. output.Count - 1 }
+        |> Seq.map (fun i -> output[i].GetTensorDataAsSpan<float32>().ToArray ())
+
 type internal AISearcher(oracle: Oracle, aiAgentTrainingMode: Option<AIAgentTrainingMode>) =
     let stepsToSwitchToAI =
         match aiAgentTrainingMode with
@@ -220,44 +225,8 @@ type internal AISearcher(oracle: Oracle, aiAgentTrainingMode: Option<AIAgentTrai
         let numOfStateAttributes = 7
         let numOfHistoryEdgeAttributes = 2
 
-        let serializeOutput (output: IDisposableReadOnlyCollection<OrtValue>) =
-            let arrayOutput =
-                seq { 0 .. output.Count - 1 }
-                |> Seq.map (fun i -> output[i].GetTensorDataAsSpan<float32>().ToArray ())
-
-            let arrayOutputJson =
-                JsonSerializer.Serialize arrayOutput
-            arrayOutputJson
-
-        let stepToString (gameState: GameState) (output: IDisposableReadOnlyCollection<OrtValue>) =
-            let gameStateJson =
-                JsonSerializer.Serialize gameState
-            let outputJson = serializeOutput output
-            let DELIM = Environment.NewLine
-            let strToSaveAsList =
-                [
-                    gameStateJson
-                    DELIM
-                    outputJson
-                    DELIM
-                ]
-            String.concat " " strToSaveAsList
 
         let createOracleRunner (pathToONNX: string, aiAgentTrainingModelOptions: Option<AIAgentTrainingModelOptions>) =
-            let stream =
-                match aiAgentTrainingModelOptions with
-                | Some options -> options.stream
-                | None -> None
-
-            let saveStep (gameState: GameState) (output: IDisposableReadOnlyCollection<OrtValue>) =
-                match stream with
-                | Some stream ->
-                    let bytes =
-                        Encoding.UTF8.GetBytes (stepToString gameState output)
-                    stream.Write (bytes, 0, bytes.Length)
-                    stream.Flush ()
-                | None -> ()
-
             let sessionOptions =
                 if useGPU then
                     SessionOptions.MakeSessionOptionWithCudaProvider (0)
@@ -471,7 +440,10 @@ type internal AISearcher(oracle: Oracle, aiAgentTrainingMode: Option<AIAgentTrai
 
                 let _ =
                     match aiAgentTrainingModelOptions with
-                    | Some _ -> saveStep gameStateOrDelta output
+                    | Some aiAgentOptions ->
+                        aiAgentOptions.stepSaver (
+                            AIGameStep (gameState = gameStateOrDelta, output = GameUtils.convertOutputToJson output)
+                        )
                     | None -> ()
 
                 stepsPlayed <- stepsPlayed + 1

diff --git a/VSharp.Explorer/Options.fs b/VSharp.Explorer/Options.fs
@@ -4,6 +4,7 @@ open System.Diagnostics
 open System.IO
 open VSharp.ML.GameServer.Messages
 open System.Net.Sockets
+open Microsoft.ML.OnnxRuntime
 
 type searchMode =
     | DFSMode
@@ -54,6 +55,17 @@ type Oracle =
 /// <param name="mapName">Name of map to play.</param>
 /// <param name="mapName">Name of map to play.</param>
 
+[<Struct>]
+type AIGameStep =
+    interface IRawOutgoingMessageBody
+    val GameState: GameState
+    val Output: seq<array<float32>>
+    new(gameState, output) =
+        {
+            GameState = gameState
+            Output = output
+        }
+
 
 type AIBaseOptions =
     {
@@ -79,7 +91,7 @@ type AIAgentTrainingModelOptions =
     {
         aiAgentTrainingOptions: AIAgentTrainingOptions
         outputDirectory: string
-        stream: Option<NetworkStream> // use it for sending steps
+        stepSaver: AIGameStep -> Unit
     }
 
 

diff --git a/VSharp.ML.GameServer.Runner/Main.fs b/VSharp.ML.GameServer.Runner/Main.fs
@@ -1,4 +1,6 @@
 open System.IO
+open System.Text
+open System.Text.Json
 open System.Net.Sockets
 open System.Reflection
 open Argu
@@ -370,20 +372,16 @@ let runTrainingSendModelMode
             oracle = None
         }
 
-    let stream =
-        let host = "localhost" // TODO: working within a local network
-        let client = new TcpClient ()
-        client.Connect (host, port)
-        client.SendBufferSize <- 2048
-        Some <| client.GetStream ()
-
+    let mutable steps = []
+    let stepSaver (aiGameStep: AIGameStep) = steps <- aiGameStep :: steps in
+    ()
     let aiOptions: AIOptions =
         Training (
             SendModel
                 {
                     aiAgentTrainingOptions = aiTrainingOptions
                     outputDirectory = outputDirectory
-                    stream = stream
+                    stepSaver = stepSaver
                 }
         )
 
@@ -410,8 +408,29 @@ let runTrainingSendModelMode
 
     printfn
         $"Running for {gameMap.MapName} finished with coverage {explorationResult.ActualCoverage}, tests {explorationResult.TestsCount}, steps {explorationResult.StepsCount},errors {explorationResult.ErrorsCount}."
-
-
+    let steps = List.rev steps
+    let stream =
+        let host = "localhost" // TODO: working within a local network
+        let client = new TcpClient ()
+        client.Connect (host, port)
+        client.SendBufferSize <- 4096
+        client.GetStream ()
+
+    let needToSendSteps =
+        let buffer = Array.zeroCreate<byte> 1
+        let bytesRead = stream.Read (buffer, 0, 1)
+        if bytesRead = 0 then
+            failwith "Connection is closed?!"
+        buffer.[0] <> byte 0
+
+    if needToSendSteps then
+        let bytes =
+            Encoding.UTF8.GetBytes (JsonSerializer.Serialize steps)
+        stream.Write (bytes, 0, bytes.Length)
+        stream.Flush ()
+        stream.Close ()
+    else
+        ()
 
 [<EntryPoint>]
 let main args =