add armhand environment

src/LyceumMuJoCo.jl

@@ -52,7 +52,7 @@ import LyceumBase: statespace,
 
                    tconstruct
 
-using LyceumBase.Tools: perturb!
+using LyceumBase.Tools: perturb!, SPoint3D
 
 export # AbstractMuJoCoEnvironment interface (an addition to AbstractEnvironment's interface)
        AbstractMuJoCoEnvironment,
@@ -138,6 +138,7 @@ end
 ####
 
 include("suite/pointmass.jl")
+include("suite/armhand/armhandpickup.jl")
 
 include("gym/swimmer-v2.jl")
 include("gym/hopper-v2.jl")

src/suite/armhand/armhandpickup.jl

@@ -0,0 +1,146 @@
+struct ArmHandPickup{S<:MJSim,O<:MultiShape} <: AbstractMuJoCoEnvironment
+    sim::S
+    obsspace::O
+    ball::Int
+    palm::Int
+    thumb::Int
+    index::Int
+    middle::Int
+    ring::Int
+    pinky::Int
+    goal::SVector{3,Float64}
+
+    function ArmHandPickup(sim::MJSim)
+        m = sim.m
+        mn, dn = sim.mn, sim.dn
+
+        obsspace = MultiShape(
+            d_thumb = ScalarShape(Float64),
+            d_index = ScalarShape(Float64),
+            d_middle = ScalarShape(Float64),
+            d_ring = ScalarShape(Float64),
+            d_pinky = ScalarShape(Float64),
+            a_thumb = ScalarShape(Float64),
+            a_index = ScalarShape(Float64),
+            a_middle = ScalarShape(Float64),
+            a_ring = ScalarShape(Float64),
+            a_pinky = ScalarShape(Float64),
+            a_close = ScalarShape(Float64),
+            handball = ScalarShape(Float64),
+            ballgoal = ScalarShape(Float64),
+            ball = VectorShape(Float64, 3),
+            palm = VectorShape(Float64, 3),
+        )
+
+        ball = jl_name2id(sim.m, MuJoCo.MJCore.mjOBJ_SITE, "ball")
+        palm = jl_name2id(sim.m, MuJoCo.MJCore.mjOBJ_SITE, "palm")
+
+        thumb = jl_name2id(sim.m, MuJoCo.MJCore.mjOBJ_SITE, "thumb_IMU")
+        index = jl_name2id(sim.m, MuJoCo.MJCore.mjOBJ_SITE, "index_IMU")
+        middle = jl_name2id(sim.m, MuJoCo.MJCore.mjOBJ_SITE, "middle_IMU")
+        ring = jl_name2id(sim.m, MuJoCo.MJCore.mjOBJ_SITE, "ring_IMU")
+        pinky = jl_name2id(sim.m, MuJoCo.MJCore.mjOBJ_SITE, "pinky_IMU")
+
+        goal = SA_F64[0.0, 0.2, 0.5]
+
+        new{typeof(sim),typeof(obsspace)}(
+            sim,
+            obsspace,
+            ball,
+            palm,
+            thumb,
+            index,
+            middle,
+            ring,
+            pinky,
+            goal,
+        )
+    end
+end
+
+ArmHandPickup() = first(tconstruct(ArmHandPickup, 1))
+
+function tconstruct(::Type{ArmHandPickup}, n::Integer)
+    modelpath = joinpath(@__DIR__, "armhand.xml")
+    Tuple(ArmHandPickup(s) for s in tconstruct(MJSim, n, modelpath, skip = 3))
+end
+
+
+@inline getsim(env::ArmHandPickup) = env.sim
+
+
+@inline obsspace(env::ArmHandPickup) = env.obsspace
+
+@propagate_inbounds function getobs!(obs, env::ArmHandPickup)
+    @boundscheck checkaxes(obsspace(env), obs)
+
+    m, d = env.sim.m, env.sim.d
+    sx = d.site_xpos
+    dmin = 0.5
+
+    ball = SPoint3D(sx, env.ball)
+    palm = SPoint3D(sx, env.palm)
+    thumb = SPoint3D(sx, env.thumb)
+    index = SPoint3D(sx, env.index)
+    middle = SPoint3D(sx, env.middle)
+    ring = SPoint3D(sx, env.ring)
+    pinky = SPoint3D(sx, env.pinky)
+    goal = ball - env.goal
+
+    @uviews obs @inbounds begin
+        shaped = obsspace(env)(obs)
+
+        shaped.ball .= ball
+        shaped.palm .= palm .- ball
+
+        shaped.handball = _sitedist(palm, ball, dmin)
+        shaped.ballgoal = _sitedist(ball, goal, dmin)
+
+        shaped.d_thumb = _sitedist(thumb, ball, dmin)
+        shaped.d_index = _sitedist(index, ball, dmin)
+        shaped.d_middle = _sitedist(middle, ball, dmin)
+        shaped.d_ring = _sitedist(ring, ball, dmin)
+        shaped.d_pinky = _sitedist(pinky, ball, dmin)
+        shaped.a_thumb = cosine_dist(goal, thumb)
+        shaped.a_index = cosine_dist(goal, index)
+        shaped.a_middle = cosine_dist(goal, middle)
+        shaped.a_ring = cosine_dist(goal, ring)
+        shaped.a_pinky = cosine_dist(goal, pinky)
+        shaped.a_close = cosine_dist(middle, thumb)
+    end
+
+    obs
+end
+
+
+@propagate_inbounds function getreward(state, action, obs, env::ArmHandPickup)
+    @boundscheck checkaxes(obsspace(env), obs)
+
+    os = obsspace(env)(obs)
+    handball = os.handball / 0.5
+    ballgoal = os.ballgoal / 0.5
+
+    reward = -handball
+    if handball < 0.06
+        reward = 2.0 - 2 * ballgoal
+        reward -= 0.1 * (os.d_thumb + os.d_index + os.d_middle + os.d_ring + os.d_pinky)
+    end
+    reward
+end
+
+@propagate_inbounds function geteval(state, action, obs, env::ArmHandPickup)
+    @boundscheck checkaxes(obsspace(env), obs)
+    obsspace(env)(obs).ball[3]
+end
+
+
+@propagate_inbounds function randreset!(rng::Random.AbstractRNG, env::ArmHandPickup)
+    fastreset_nofwd!(env.sim)
+    env.sim.d.qpos[1] = rand(rng, Uniform(-0.15, 0.15))
+    env.sim.d.qpos[2] = rand(rng, Uniform(-0.1, 0.1))
+    forward!(env.sim)
+    env
+end
+
+
+@inline _sitedist(s1, s2, dmin) = min(euclidean(s1, s2), dmin)

test/runtests.jl

@@ -3,7 +3,8 @@ using Test, Random, Shapes, Pkg, LinearAlgebra, BenchmarkTools
 using MuJoCo: TESTMODELXML
 
 const LYCEUM_SUITE = [
-    (LyceumMuJoCo.PointMass, (), ())
+    (LyceumMuJoCo.PointMass, (), ()),
+    (LyceumMuJoCo.ArmHandPickup, (), ()),
 ]
 
 const GYM = [
@@ -32,4 +33,4 @@ end
         @testset "DMC" begin test_group(DMC) end
     end
 
-end
+end