From 7910686acc8b02668ded8ce54a2c38f5d288aa22 Mon Sep 17 00:00:00 2001
From: Oldrich Krupauer <oldrich.krupauer@pm.me>
Date: Tue, 22 Oct 2024 08:18:30 +0200
Subject: [PATCH] Added option to change appearance of relfection from python

---
 python_trainer/environment_parameters.py      | 32 +++++++++++++++++++
 python_trainer/train.py                       | 14 ++------
 .../Assets/Materials/Road/Amazon road.mat     |  2 +-
 unity_env/Assets/Scenes/Main.unity            |  3 ++
 .../Assets/Scripts/TrainingReplicator.cs      | 18 +++++++++++
 5 files changed, 57 insertions(+), 12 deletions(-)
 create mode 100644 python_trainer/environment_parameters.py

diff --git a/python_trainer/environment_parameters.py b/python_trainer/environment_parameters.py
new file mode 100644
index 00000000..94a57b0e
--- /dev/null
+++ b/python_trainer/environment_parameters.py
@@ -0,0 +1,32 @@
+def set_parameters(data_channel) -> None:
+    # Wide - 15
+    # Slim - 10
+    data_channel.set_int_parameter("roadSize", 15)
+
+    # 0 -> Amazon road
+    # 1 -> Black & white road
+    data_channel.set_int_parameter("roadColor", 0)
+
+    data_channel.set_bool_parameter("randomBackgroundColor", True)
+
+    data_channel.set_float_parameter("changingBackgroundColorSpeed", 0.75)
+
+    # When the parameter 'backgroundColor' is not set, it is going to generate random colors
+    # Values of the channels are <0, 255>
+    # data_channel.set_color_parameter('backgroundColor', (255, 0, 0))
+
+    # How well you can see the reflection
+    # Default: 1.0
+    data_channel.set_float_parameter("reflectionStrength", 1.0)
+
+    # How scaled are the reflections in the X and Y directions
+    # Scaling both the numbers up will make the reflections smaller
+    # Default: 0.4
+    data_channel.set_float_parameter("noiseScaleX", 0.4)
+    # Default: 0.05
+    data_channel.set_float_parameter("noiseScaleY", 0.05)
+
+    # The animation speed of the reflection
+    # Range <0, 1>
+    # Default: 0.4
+    data_channel.set_float_parameter("noiseSpeed", 0.4)
diff --git a/python_trainer/train.py b/python_trainer/train.py
index cd6c51b3..03016c9c 100644
--- a/python_trainer/train.py
+++ b/python_trainer/train.py
@@ -5,6 +5,7 @@
 
 import torch
 from data_channel import DataChannel
+from environment_parameters import set_parameters
 from keyboard_listener import KeyboardListener
 from mlagents_envs.environment import UnityEnvironment
 from mlagents_envs.side_channel.engine_configuration_channel import EngineConfigurationChannel
@@ -86,17 +87,8 @@ def print_env_info(env: UnityEnvironment) -> None:
         side_channels=[engine_channel, data_channel],
     )
 
-    # Wide - 15
-    # Slim - 10
-    data_channel.set_int_parameter("roadSize", 15)
-    # 0 -> Amazon road
-    # 1 -> Black & white road
-    data_channel.set_int_parameter("roadColor", 0)
-    data_channel.set_bool_parameter("randomBackgroundColor", True)
-    data_channel.set_float_parameter("changingBackgroundColorSpeed", 0.75)
-    # When the parameter 'backgroundColor' is not set, the default color is the color of the amazon color
-    # Values of the channels are <0, 255>
-    # data_channel.set_color_parameter('backgroundColor', (255, 0, 0))
+    set_parameters(data_channel)
+
     engine_channel.set_configuration_parameters(time_scale=TIME_SCALE)
     env.reset()
 
diff --git a/unity_env/Assets/Materials/Road/Amazon road.mat b/unity_env/Assets/Materials/Road/Amazon road.mat
index 6a665bde..d7a3de60 100644
--- a/unity_env/Assets/Materials/Road/Amazon road.mat	
+++ b/unity_env/Assets/Materials/Road/Amazon road.mat	
@@ -85,7 +85,7 @@ Material:
     - _NoiseScaleY: 0.05
     - _OcclusionStrength: 1
     - _Parallax: 0.02
-    - _ReflectionStrength: 1
+    - _ReflectionStrength: 0
     - _SetTexture: 1
     - _Smoothness: 0
     - _SmoothnessTextureChannel: 1
diff --git a/unity_env/Assets/Scenes/Main.unity b/unity_env/Assets/Scenes/Main.unity
index 7d5e0700..2ac02771 100644
--- a/unity_env/Assets/Scenes/Main.unity
+++ b/unity_env/Assets/Scenes/Main.unity
@@ -723,6 +723,9 @@ MonoBehaviour:
   roadColor: 1
   backgroundColor: {r: 0, g: 0, b: 0, a: 0}
   randomBackgroundColor: 0
+  roadMaterials:
+  - {fileID: 2100000, guid: d63a402699f7f904c8bad171fdc44dca, type: 2}
+  - {fileID: 2100000, guid: 39a14f8831164d0b6b9637adb0cdbddf, type: 2}
 --- !u!1001 &7475695033151031654
 PrefabInstance:
   m_ObjectHideFlags: 0
diff --git a/unity_env/Assets/Scripts/TrainingReplicator.cs b/unity_env/Assets/Scripts/TrainingReplicator.cs
index 3272b04c..13248d8d 100644
--- a/unity_env/Assets/Scripts/TrainingReplicator.cs
+++ b/unity_env/Assets/Scripts/TrainingReplicator.cs
@@ -28,6 +28,8 @@ public enum RoadColor
         private float changingColorSpeed;
         private bool displayRandomColorInMain = false;
 
+        public Material[] roadMaterials;
+
         struct HSV
         {
             public float hue;
@@ -71,6 +73,8 @@ void Start()
                 );
             }
 
+            SetMaterial();
+
             AddAreas();
 
             carCameras = GameObject.FindGameObjectsWithTag("CarCamera");
@@ -105,6 +109,20 @@ private void ChangeCameraBackgroundColor()
                 GameObject.Find("Camera").GetComponent<Camera>().backgroundColor = backgroundColor;
         }
 
+        void SetMaterial()
+        {
+            foreach (var mat in roadMaterials)
+            {
+                mat.SetFloat(
+                    "_ReflectionStrength",
+                    DataChannel.getParameter("reflectionStrength", 1f)
+                );
+                mat.SetFloat("_NoiseScaleX", DataChannel.getParameter("noiseScaleX", 0.4f));
+                mat.SetFloat("_NoiseScaleY", DataChannel.getParameter("noiseScaleY", 0.05f));
+                mat.SetFloat("_Speed", DataChannel.getParameter("noiseSpeed", 0.3f));
+            }
+        }
+
         private void AddAreas()
         {
             for (int i = 0; i < numAreas; i++)