File size: 5,170 Bytes
05c9ac2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
#if UNITY_2020_1_OR_NEWER
using UnityEngine;
using NUnit.Framework;
using Unity.MLAgents.Extensions.Sensors;


namespace Unity.MLAgents.Extensions.Tests.Sensors
{
    public class ArticulationBodySensorTests
    {
        [Test]
        public void TestNullRootBody()
        {
            var gameObj = new GameObject();

            var sensorComponent = gameObj.AddComponent<ArticulationBodySensorComponent>();
            var sensor = sensorComponent.CreateSensors()[0];
            SensorTestHelper.CompareObservation(sensor, new float[0]);
        }

        [Test]
        public void TestSingleBody()
        {
            var gameObj = new GameObject();
            var articulationBody = gameObj.AddComponent<ArticulationBody>();
            var sensorComponent = gameObj.AddComponent<ArticulationBodySensorComponent>();
            sensorComponent.RootBody = articulationBody;
            sensorComponent.Settings = new PhysicsSensorSettings
            {
                UseModelSpaceLinearVelocity = true,
                UseLocalSpaceTranslations = true,
                UseLocalSpaceRotations = true
            };

            var sensor = sensorComponent.CreateSensors()[0];
            sensor.Update();
            var expected = new[]
            {
                0f, 0f, 0f, // ModelSpaceLinearVelocity
                0f, 0f, 0f, // LocalSpaceTranslations
                0f, 0f, 0f, 1f // LocalSpaceRotations
            };
            SensorTestHelper.CompareObservation(sensor, expected);
            Assert.AreEqual(expected.Length, sensor.GetObservationSpec().Shape[0]);
        }

        [Test]
        public void TestBodiesWithJoint()
        {
            var rootObj = new GameObject();
            var rootArticBody = rootObj.AddComponent<ArticulationBody>();

            var middleGamObj = new GameObject();
            var middleArticBody = middleGamObj.AddComponent<ArticulationBody>();
            middleArticBody.AddForce(new Vector3(0f, 1f, 0f));
            middleGamObj.transform.SetParent(rootObj.transform);
            middleGamObj.transform.localPosition = new Vector3(13.37f, 0f, 0f);
            middleArticBody.jointType = ArticulationJointType.RevoluteJoint;

            var leafGameObj = new GameObject();
            var leafArticBody = leafGameObj.AddComponent<ArticulationBody>();
            leafGameObj.transform.SetParent(middleGamObj.transform);
            leafGameObj.transform.localPosition = new Vector3(4.2f, 0f, 0f);
            leafArticBody.jointType = ArticulationJointType.PrismaticJoint;
            leafArticBody.linearLockZ = ArticulationDofLock.LimitedMotion;
            leafArticBody.zDrive = new ArticulationDrive
            {
                lowerLimit = -3,
                upperLimit = 1
            };


#if UNITY_2020_2_OR_NEWER
            // ArticulationBody.velocity is read-only in 2020.1
            rootArticBody.velocity = new Vector3(1f, 0f, 0f);
            middleArticBody.velocity = new Vector3(0f, 1f, 0f);
            leafArticBody.velocity = new Vector3(0f, 0f, 1f);
#endif

            var sensorComponent = rootObj.AddComponent<ArticulationBodySensorComponent>();
            sensorComponent.RootBody = rootArticBody;
            sensorComponent.Settings = new PhysicsSensorSettings
            {
                UseModelSpaceTranslations = true,
                UseLocalSpaceTranslations = true,
#if UNITY_2020_2_OR_NEWER
                UseLocalSpaceLinearVelocity = true
#endif
            };

            var sensor = sensorComponent.CreateSensors()[0];
            sensor.Update();
            var expected = new[]
            {
                // Model space
                0f, 0f, 0f, // Root pos
                13.37f, 0f, 0f, // Middle pos
                leafGameObj.transform.position.x, 0f, 0f, // Leaf pos

                // Local space
                0f, 0f, 0f, // Root pos
                13.37f, 0f, 0f, // Attached pos
                4.2f, 0f, 0f, // Leaf pos

#if UNITY_2020_2_OR_NEWER
                0f, 0f, 0f, // Root vel
                -1f, 1f, 0f, // Attached vel
                0f, -1f, 1f // Leaf vel
#endif
            };
            SensorTestHelper.CompareObservation(sensor, expected);
            Assert.AreEqual(expected.Length, sensor.GetObservationSpec().Shape[0]);

            // Update the settings to only process joint observations
            sensorComponent.Settings = new PhysicsSensorSettings
            {
                UseJointForces = true,
                UseJointPositionsAndAngles = true,
            };

            sensor = sensorComponent.CreateSensors()[0];
            sensor.Update();

            expected = new[]
            {
                // revolute
                0f, 1f, // joint1.position (sin and cos)
                0f, // joint1.force

                // prismatic
                0.5f, // joint2.position (interpolate between limits)
                0f, // joint2.force
            };
            SensorTestHelper.CompareObservation(sensor, expected);
            Assert.AreEqual(expected.Length, sensor.GetObservationSpec().Shape[0]);
        }
    }
}
#endif // #if UNITY_2020_1_OR_NEWER