From eb8389feaddf6c442becdbd0fc25cfb6cac845fd Mon Sep 17 00:00:00 2001
From: Tom Erez <etom@google.com>
Date: Fri, 30 Dec 2022 18:37:17 -0800
Subject: [PATCH] Find MjScene at OnEnable: ensure MuJoCo compilation happens
 only once at startup.

PiperOrigin-RevId: 498655077
Change-Id: Ibed8f36a880426063b7729fc66b15c67adcc9a65
---
 unity/Runtime/Components/Joints/MjBaseJoint.cs | 3 +--
 unity/Runtime/Components/MjComponent.cs        | 2 +-
 unity/Runtime/Components/MjScene.cs            | 1 +
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/unity/Runtime/Components/Joints/MjBaseJoint.cs b/unity/Runtime/Components/Joints/MjBaseJoint.cs
index 55fe0b2f..c5ad6241 100644
--- a/unity/Runtime/Components/Joints/MjBaseJoint.cs
+++ b/unity/Runtime/Components/Joints/MjBaseJoint.cs
@@ -49,8 +49,7 @@ namespace Mujoco {
       DofAddress = model->jnt_dofadr[MujocoId];
     }
 
-    protected override void Start() {
-      base.Start();
+    protected void Start() {
       MjBaseBody grandparentBody, parentBody;
       GetConnectedBodies(out grandparentBody, out parentBody);
       GrandParentBody = grandparentBody;
diff --git a/unity/Runtime/Components/MjComponent.cs b/unity/Runtime/Components/MjComponent.cs
index 3b69668f..ae6f8094 100644
--- a/unity/Runtime/Components/MjComponent.cs
+++ b/unity/Runtime/Components/MjComponent.cs
@@ -77,7 +77,7 @@ public abstract class MjComponent : MonoBehaviour {
 
   private bool _sceneExcludesMe = false;
 
-  protected unsafe virtual void Start() {
+  protected unsafe virtual void OnEnable() {
     if (MjScene.Instance == null) {
       throw new Exception("MuJoCo Scene not found");
     }
diff --git a/unity/Runtime/Components/MjScene.cs b/unity/Runtime/Components/MjScene.cs
index 77f4c654..9e8eed7f 100644
--- a/unity/Runtime/Components/MjScene.cs
+++ b/unity/Runtime/Components/MjScene.cs
@@ -85,6 +85,7 @@ public class MjScene : MonoBehaviour {
   public EventHandler<MjStepArgs> preDestroyEvent;
 
   protected unsafe void Start() {
+    SceneRecreationAtLateUpdateRequested = false;
     CreateScene();
   }