From 76e020362740865bec487d745f629c720e8f5618 Mon Sep 17 00:00:00 2001
From: Eli VanderBilt <eliv@allenai.org>
Date: Tue, 16 Jul 2024 10:25:24 -0700
Subject: [PATCH] Added post-processing to stretch-robot initialization

---
 .../RoboTHOR_Post_Profile.asset               | 37 ++++++++++++++++++-
 .../Assets/Scripts/StretchAgentController.cs  | 25 ++++++++++++-
 2 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/unity/Assets/Physics/SimObjsPhysics/RoboTHOR Objects/RoboTHOR_Assets_Environment/RoboTHOR_Post_Profile.asset b/unity/Assets/Physics/SimObjsPhysics/RoboTHOR Objects/RoboTHOR_Assets_Environment/RoboTHOR_Post_Profile.asset
index 34329e9a03..6d14d68c9c 100644
--- a/unity/Assets/Physics/SimObjsPhysics/RoboTHOR Objects/RoboTHOR_Assets_Environment/RoboTHOR_Post_Profile.asset	
+++ b/unity/Assets/Physics/SimObjsPhysics/RoboTHOR Objects/RoboTHOR_Assets_Environment/RoboTHOR_Post_Profile.asset	
@@ -16,6 +16,7 @@ MonoBehaviour:
   - {fileID: 114156960648436490}
   - {fileID: 114125738884031076}
   - {fileID: 114774809263292024}
+  - {fileID: 8332371402478750980}
 --- !u!114 &114125738884031076
 MonoBehaviour:
   m_ObjectHideFlags: 3
@@ -75,7 +76,7 @@ MonoBehaviour:
   m_Script: {fileID: 11500000, guid: c1cb7e9e120078f43bce4f0b1be547a7, type: 3}
   m_Name: AmbientOcclusion
   m_EditorClassIdentifier: 
-  active: 0
+  active: 1
   enabled:
     overrideState: 1
     value: 1
@@ -1423,3 +1424,37 @@ MonoBehaviour:
       - 0.5
       - 0.5
       - 0.5
+--- !u!114 &8332371402478750980
+MonoBehaviour:
+  m_ObjectHideFlags: 3
+  m_CorrespondingSourceObject: {fileID: 0}
+  m_PrefabInstance: {fileID: 0}
+  m_PrefabAsset: {fileID: 0}
+  m_GameObject: {fileID: 0}
+  m_Enabled: 1
+  m_EditorHideFlags: 0
+  m_Script: {fileID: 11500000, guid: 9b77c5407dc277943b591ade9e6b18c5, type: 3}
+  m_Name: LensDistortion
+  m_EditorClassIdentifier: 
+  active: 0
+  enabled:
+    overrideState: 1
+    value: 1
+  intensity:
+    overrideState: 1
+    value: 35
+  intensityX:
+    overrideState: 0
+    value: 1
+  intensityY:
+    overrideState: 0
+    value: 1
+  centerX:
+    overrideState: 0
+    value: 0
+  centerY:
+    overrideState: 0
+    value: 0
+  scale:
+    overrideState: 1
+    value: 1.03
diff --git a/unity/Assets/Scripts/StretchAgentController.cs b/unity/Assets/Scripts/StretchAgentController.cs
index b28e9de058..e737218795 100644
--- a/unity/Assets/Scripts/StretchAgentController.cs
+++ b/unity/Assets/Scripts/StretchAgentController.cs
@@ -29,6 +29,7 @@ public partial class StretchAgentController : ArmAgentController {
         );
         private Vector3 defaultSecondaryCameraLocalRotation = new Vector3(50f, 90f, 0);
         private float defaultSecondaryCameraFieldOfView = 59f;
+        private bool distortLens = true;
 
         public StretchAgentController(
             BaseAgentComponent baseAgentComponent,
@@ -57,8 +58,6 @@ public override ActionFinished InitializeBody(ServerAction initializeAction) {
             cc.center = m_CharacterController.center;
             cc.radius = m_CharacterController.radius;
             cc.height = m_CharacterController.height;
-            m_Camera.GetComponent<PostProcessVolume>().enabled = true;
-            m_Camera.GetComponent<PostProcessLayer>().enabled = true;
 
             // set camera stand/crouch local positions for Tall mode
             standingLocalCameraPosition = m_Camera.transform.localPosition;
@@ -92,6 +91,12 @@ public override ActionFinished InitializeBody(ServerAction initializeAction) {
             fp_camera_2.transform.localEulerAngles = defaultSecondaryCameraLocalRotation;
             fp_camera_2.fieldOfView = defaultSecondaryCameraFieldOfView;
 
+            // set up lens distortion for stretch bot cameras
+            if (distortLens == true) {
+                enableLensDistortion(m_Camera);
+                enableLensDistortion(fp_camera_2);
+            }
+
             // limit camera from looking too far down/up
             if (Mathf.Approximately(initializeAction.maxUpwardLookAngle, 0.0f)) {
                 this.maxUpwardLookAngle = 25f;
@@ -133,6 +138,22 @@ out secondaryCameraParams
             return ActionFinished.Success;
         }
 
+        public void enableLensDistortion(Camera camera) {
+            // "enableLensDistortion" - Set this up as an external function that takes a camera, then run it for both!!!
+            // don't forget to emulate this logic for LocoBot, but only if you think it's worth standardizing, and not one-offing like here...
+            PostProcessVolume cameraPPVolume = camera.GetComponent<PostProcessVolume>();
+            cameraPPVolume.enabled = true;
+        
+            PostProcessLayer cameraPPLayer = camera.GetComponent<PostProcessLayer>();
+            cameraPPLayer.enabled = true;
+
+            PostProcessProfile cameraPPProfile = cameraPPVolume.profile;
+
+            if (cameraPPProfile.TryGetSettings(out LensDistortion lensDistortion)) {
+                lensDistortion.active = true; // Toggle Lens Distortion on based on some condition
+            }
+        }
+
         private ArmController getArmImplementation() {
             Stretch_Robot_Arm_Controller arm =
                 GetComponentInChildren<Stretch_Robot_Arm_Controller>();