Vmas integration (#31)

* feat: Setup module for the VMAS integration

* feat: completed integration with vmas

* feat: add rendering and cuda support

* feat: add wandb logging and generalize logger for CTDESystem, Learner and DTDEAgent

* feat: refactor classes names and deprecate VmasCTDEAgent class

* chore: delete unnecessary class

* feat: add CohesionAndCollision scenario without LiDAR

* feat: boost performances of CohesionAndCollisionNoLidar scenario

* feat: boost performances of CohesionAndCollisionNoLidar scenario

* feat: add snapshot for tests

* feat: add DSL for defining reward functions

---------

Co-authored-by: Filippo Cavallari <[email protected]>

Author: domm99

138-2024-01-16-18-24-46-agent-0

@@ -18,4 +18,5 @@ object AutodiffDevice {
   def apply() =
     deepLearningLib()
       .device(if (deepLearningLib().cuda.is_available().as[Boolean]) "cuda" else "cpu")
+//        .device("cpu")
 }

scarlib-core/src/main/scala/it/unibo/scarlib/core/model/AutodiffDevice.scala

@@ -10,12 +10,13 @@
 package it.unibo.scarlib.core.model
 
 import it.unibo.scarlib.core.neuralnetwork.{NeuralNetworkEncoding, SimpleSequentialDQN, TorchSupport}
-import it.unibo.scarlib.core.util.TorchLiveLogger
+import it.unibo.scarlib.core.util.{Logger, TorchLiveLogger}
 import me.shadaj.scalapy.py
 import me.shadaj.scalapy.py.{PyQuote, SeqConverters}
 
 import java.text.SimpleDateFormat
 import java.util.Date
+import scala.reflect.io.{File, Path}
 import scala.util.Random
 
 /** The DQN learning algorithm
@@ -27,7 +28,8 @@ import scala.util.Random
 class DeepQLearner(
     memory: ReplayBuffer[State, Action],
     actionSpace: Seq[Action],
-    learningConfiguration: LearningConfiguration
+    learningConfiguration: LearningConfiguration,
+    logger: Logger
 )(implicit encoding: NeuralNetworkEncoding[State]) extends Learner {
 
   private val random = learningConfiguration.random
@@ -71,7 +73,7 @@ class DeepQLearner(
       val expectedValue = (nextStateValues * gamma) + rewards
       val criterion = TorchSupport.neuralNetworkModule().SmoothL1Loss()
       val loss = criterion(stateActionValue, expectedValue.unsqueeze(1))
-      TorchLiveLogger.logScalar("Loss", loss.item().as[Double], updates)
+      logger.logScalar("Loss", loss.item().as[Double], updates)
       optimizer.zero_grad()
       loss.backward()
       it.unibo.scarlib.core.neuralnetwork.TorchSupport
@@ -93,9 +95,14 @@ class DeepQLearner(
       .deepLearningLib()
       .save(
         targetNetwork.state_dict(),
-        s"${learningConfiguration.snapshotPath}-$episode-$timeMark-agent-$agentId"
+        s"${learningConfiguration.snapshotPath}${File.separator}$episode-$timeMark-agent-$agentId"
       )
   }
+
+  override def loadSnapshot(path: String): Unit = {
+    targetNetwork.load_state_dict(TorchSupport.deepLearningLib().load(path, map_location = AutodiffDevice()))
+    policyNetwork.load_state_dict(TorchSupport.deepLearningLib().load(path, map_location = AutodiffDevice()))
+  }
 }
 
 object DeepQLearner {

scarlib-core/src/main/scala/it/unibo/scarlib/core/model/DeepQLearner.scala

@@ -23,4 +23,5 @@ trait Learner {
   /** Takes a snapshot of the current policy */
   def snapshot(episode: Int, agentId: Int): Unit
 
+  def loadSnapshot(path: String): Unit
 }

scarlib-core/src/main/scala/it/unibo/scarlib/core/model/Learner.scala

@@ -18,5 +18,7 @@ trait DeepLearningSupport[M]{
 
   def logger(): M
 
+  def arrayModule: M
+
 }
 

scarlib-core/src/main/scala/it/unibo/scarlib/core/neuralnetwork/DeepLearningSupport.scala

@@ -20,4 +20,6 @@ object TorchSupport extends DeepLearningSupport[py.Module] {
   override def optimizerModule(): py.Module = py.module("torch.optim")
 
   override def logger(): py.Module = py.module("torch.utils.tensorboard")
+
+  override def arrayModule: py.Module = py.module("numpy")
 }

scarlib-core/src/main/scala/it/unibo/scarlib/core/neuralnetwork/TorchSupport.scala

@@ -9,7 +9,8 @@
 
 package it.unibo.scarlib.core.system
 
-import it.unibo.scarlib.core.model._
+import it.unibo.scarlib.core.model.{Action, Decay, DeepQLearner, Environment, LearningConfiguration, ReplayBuffer, State}
+import it.unibo.scarlib.core.util.{Logger, TorchLiveLogger}
 import it.unibo.scarlib.core.neuralnetwork.{NeuralNetworkEncoding, NeuralNetworkSnapshot}
 
 import scala.annotation.tailrec
@@ -30,12 +31,16 @@ class CTDESystem(
                   environment: Environment,
                   dataset: ReplayBuffer[State, Action],
                   actionSpace: Seq[Action],
-                  learningConfiguration: LearningConfiguration
+                  learningConfiguration: LearningConfiguration,
+                  logger: Logger = TorchLiveLogger
 )(implicit context: ExecutionContext, encoding: NeuralNetworkEncoding[State]) {
 
   private val epsilon: Decay[Double] = learningConfiguration.epsilon
-  private val learner =
-    new DeepQLearner(dataset, actionSpace, learningConfiguration)
+
+  private val learner = new DeepQLearner(dataset, actionSpace, learningConfiguration, logger)
+
+
+
 
   /** Starts the learning process
    *
@@ -47,6 +52,7 @@ class CTDESystem(
     @tailrec
     def singleEpisode(time: Int): Unit =
       if (time > 0) {
+        println("Time: " + time)
         agents.foreach(_.notifyNewPolicy(learner.behavioural))
         Await.ready(Future.sequence(agents.map(_.step())), scala.concurrent.duration.Duration.Inf)
         environment.log()
@@ -65,6 +71,11 @@ class CTDESystem(
 
   }
 
+    final def learn(episodes: Int, episodeLength: Int, snapshot: String): Unit = {
+        learner.loadSnapshot(snapshot)
+        learn(episodes, episodeLength)
+    }
+
   /** Starts the testing process
    *
    * @param episodeLength the length of the episode
@@ -79,6 +90,7 @@ class CTDESystem(
 
     @tailrec
     def episode(time: Int): Unit = {
+      println(time)
       if (time > 0) {
         Await.ready(Future.sequence(agents.map(_.step())), scala.concurrent.duration.Duration.Inf)
         episode(time - 1)

scarlib-core/src/main/scala/it/unibo/scarlib/core/system/CTDESystem.scala

@@ -9,8 +9,9 @@
 
 package it.unibo.scarlib.core.system
 
-import it.unibo.scarlib.core.model._
+import it.unibo.scarlib.core.model.{Action, Agent, AgentMode, Decay, DeepQLearner, Environment, Experience, LearningConfiguration, ReplayBuffer, State}
 import it.unibo.scarlib.core.neuralnetwork.{NeuralNetworkEncoding, NeuralNetworkSnapshot}
+import it.unibo.scarlib.core.util.{Logger, TorchLiveLogger}
 
 import scala.reflect.io.File
 import scala.concurrent.ExecutionContext.Implicits.global
@@ -31,12 +32,13 @@ class DTDEAgent(
                           actionSpace: Seq[Action],
                           datasetSize: Int,
                           agentMode: AgentMode = AgentMode.Training,
-                          learningConfiguration: LearningConfiguration
+                          learningConfiguration: LearningConfiguration,
+                          logger: Logger = TorchLiveLogger
 )(implicit encoding: NeuralNetworkEncoding[State]) extends Agent {
 
   private val dataset: ReplayBuffer[State, Action] = ReplayBuffer[State, Action](datasetSize)
   private val epsilon: Decay[Double] = learningConfiguration.epsilon
-  private val learner = new DeepQLearner(dataset, actionSpace, learningConfiguration)
+  private val learner = new DeepQLearner(dataset, actionSpace, learningConfiguration, logger)
   private var testPolicy: State => Action = _
 
   /** A single interaction of the agent with the environment */

scarlib-core/src/main/scala/it/unibo/scarlib/core/system/DTDEAgent.scala

@@ -0,0 +1,9 @@
+package it.unibo.scarlib.core.util
+
+import me.shadaj.scalapy.py
+
+trait Logger {
+    def logScalar(tag: String, value: Double, tick: Int): Unit
+
+    def logAny(tag: String, value: py.Dynamic, tick: Int): Unit
+}

scarlib-core/src/main/scala/it/unibo/scarlib/core/util/Logger.scala

@@ -12,11 +12,12 @@ package it.unibo.scarlib.core.util
 import it.unibo.scarlib.core.neuralnetwork.TorchSupport
 import me.shadaj.scalapy.py
 
-object TorchLiveLogger {
+object TorchLiveLogger extends Logger{
     private val writer = TorchSupport.logger().SummaryWriter()
 
     def logScalar(tag: String, value: Double, tick: Int): Unit = writer.add_scalar(tag, value, tick)
 
     def logAny(tag: String, value: py.Dynamic, tick: Int): Unit = writer.add_scalar(tag, value, tick)
 }
 
+