Merge pull request #12 from rejuvyesh/jkg/envtest

janbruedigam · web-flow · commit e96c658287c0 · 2022-03-11T08:21:23.000+01:00
Add some basic tests for env interface
diff --git a/environments/ant/methods/env.jl b/environments/ant/methods/env.jl
@@ -41,7 +41,7 @@ function ant(;
         nx = maximal_dimension(mechanism)
     end
     nu = 8
-    no = nx
+    no = nx + length(mechanism.contacts)
 
     aspace = BoxSpace(nu, 
         low=(-ones(nu)), 
diff --git a/environments/box/methods/env.jl b/environments/box/methods/env.jl
@@ -22,15 +22,15 @@ function block(;
     opts_grad=SolverOptions(rtol=3.0e-4, btol=3.0e-4, undercut=1.5),
     T=Float64)
 
-    mechanism = get_box(
+    mechanism = get_block(
         timestep=timestep, 
         gravity=gravity, 
         friction_coefficient=friction_coefficient, 
         side=side, 
         contact=contact, 
         contact_type=contact_type)
 
-    initialize_box!(mechanism)
+    initialize_block!(mechanism)
 
     if representation == :minimal
         nx = minimal_dimension(mechanism)
diff --git a/environments/box/methods/initialize.jl b/environments/box/methods/initialize.jl
@@ -1,4 +1,4 @@
-function get_box(; 
+function get_block(; 
     timestep=0.01, 
     gravity=[0.0; 0.0; -9.81],
     friction_coefficient=0.8, 
@@ -58,7 +58,7 @@ function get_box(;
     return mech
 end
 
-function initialize_box!(mechanism::Mechanism{T};
+function initialize_block!(mechanism::Mechanism{T};
         x=[0.0, 0.0, 1.0],
         q=Quaternion(1.0, 0.0, 0.0, 0.0),
         v=[1.0, 0.3, 0.2],
diff --git a/environments/hopper/methods/env.jl b/environments/hopper/methods/env.jl
@@ -39,7 +39,7 @@ function hopper(;
         nx = maximal_dimension(mechanism)
     end
     nu = 3
-    no = nx
+    no = nx - 1 # full_state is not being returned by default 
 
     # values taken from Mujoco's model, combining the control range -1, 1 and the motor gears.
     aspace = BoxSpace(nu, 
@@ -140,7 +140,7 @@ function cost(env::Environment{Hopper}, x, u;
     return c
 end
 
-function is_done(::Environment{Hopper}, x)
+function is_done(env::Environment{Hopper}, x)
     nx = minimal_dimension(env.mechanism)
     if env.representation == :minimal
         x0 = x
diff --git a/environments/rexhopper/methods/env.jl b/environments/rexhopper/methods/env.jl
@@ -93,7 +93,7 @@ function Base.reset(env::Environment{RexHopper};
     else
         # initialize above the ground to make sure that with random initialization we do not violate the ground constraint.
         initialize!(env.mechanism, :rexhopper)
-        x0 = get_minimal_state(env.mechanism)
+        x = get_minimal_state(env.mechanism)
         nx = minimal_dimension(env.mechanism)
         z = minimal_to_maximal(env.mechanism, x)
         set_maximal_state!(env.mechanism, z)
diff --git a/environments/walker/methods/env.jl b/environments/walker/methods/env.jl
@@ -39,7 +39,7 @@ function walker(;
         nx = maximal_dimension(mechanism)
     end
     nu = 6
-    no = nx
+    no = nx-1 # full_state is false by default
 
     # values taken from Mujoco's model, combining the control range -1, 1 and the motor gears.
     aspace = BoxSpace(nu, 
@@ -136,7 +136,7 @@ function cost(env::Environment{Walker}, x, u;
     return c
 end
 
-function is_done(::Environment{Walker}, x)
+function is_done(env::Environment{Walker}, x)
     nx = minimal_dimension(env.mechanism)
     if env.representation == :minimal
         x0 = x
diff --git a/examples/simulation/cone_compare.jl b/examples/simulation/cone_compare.jl
@@ -32,7 +32,7 @@ opts = SolverOptions(rtol=1.0e-6, btol=1.0e-6)
 
 # ## Linear cone
 color_lc = orange;
-mech_lc = get_mechanism(:box,
+mech_lc = get_mechanism(:block,
     timestep=timestep,
     gravity=gravity,
     friction_coefficient=friction_coefficient,
@@ -41,7 +41,7 @@ mech_lc = get_mechanism(:box,
     color=color_lc);
 
 # ## Simulate
-initialize!(mech_lc, :box,
+initialize!(mech_lc, :block,
     x=x0,
     q=one(Quaternion),
     v=v0,
@@ -59,7 +59,7 @@ via, anim = visualize(mech_lc, storage_lc,
 
 # ## Nonlinear cone
 color_nc = cyan;
-mech_nc = get_mechanism(:box,
+mech_nc = get_mechanism(:block,
     timestep=timestep,
     gravity=gravity,
     friction_coefficient=friction_coefficient,
@@ -68,7 +68,7 @@ mech_nc = get_mechanism(:box,
     color=color_nc);
 
 # ## Simulate
-initialize!(mech_nc, :box,
+initialize!(mech_nc, :block,
     x=x0,
     q=one(Quaternion),
     v=v0,
@@ -91,15 +91,15 @@ vis, anim = visualize(mech_nc, storage_nc,
 # ## MuJoCo pyramidal cone
 color_mjlc = magenta;
 
-mech_mjlc = get_mechanism(:box,
+mech_mjlc = get_mechanism(:block,
     timestep=timestep,
     gravity=gravity,
     friction_coefficient=friction_coefficient,
     contact_type=:linear,
     mode=:box, color=color_mjlc);
 
 # ## Load
-initialize!(mech_mjlc, :box,
+initialize!(mech_mjlc, :block,
     x=x0,
     q=one(Quaternion),
     v=v0,
@@ -125,15 +125,15 @@ setobject!(vis[:path_mjlc], MeshCat.Line(points_mjlc, line_mat_mjlc))
 # ## MuJoCo elliptic cone
 color_mjnc = RGBA(0,0,0);
 
-mech_mjnc = get_mechanism(:box,
+mech_mjnc = get_mechanism(:block,
     timestep=timestep,
     gravity=gravity,
     friction_coefficient=friction_coefficient,
     contact_type=:linear,
     mode=:box, color=color_mjnc);
 
 # ## Load
-initialize!(mech_mjnc, :box,
+initialize!(mech_mjnc, :block,
     x=x0,
     q=one(Quaternion),
     v=v0,
diff --git a/examples/system_identification/data/data.jl b/examples/system_identification/data/data.jl
@@ -64,7 +64,7 @@ function generate_hardware_dataset(;N::Int=10,
 	timestep= 1/148 * S
 	gscaled = -9.81*20
 
-    mechanism = get_mechanism(:box, timestep=timestep, gravity=gravityscaled);
+    mechanism = get_mechanism(:block, timestep=timestep, gravity=gravityscaled);
     trajs = []
 	pairs = []
     for i = 1:N
diff --git a/examples/system_identification/hardware_examples/visualization.jl b/examples/system_identification/hardware_examples/visualization.jl
@@ -72,7 +72,7 @@ function d2data(d)
 			]
 	return data
 end
-mech = get_mechanism(:box, timestep=timestep/S, gravity=gravityscaled, friction_coefficient=Dsol[end][1], radius=0.00, side=2.0, mode=:box);
+mech = get_mechanism(:block, timestep=timestep/S, gravity=gravityscaled, friction_coefficient=Dsol[end][1], radius=0.00, side=2.0, mode=:box);
 set_simulator_data!(mech, d2data(Dsol[end]))
 id = 7#4,6,7,8
 traj_truth = trajs1[id]
@@ -81,7 +81,7 @@ v15 = traj_truth.v[1][1]
 q2 = traj_truth.q[1][1]
 ϕ15 = traj_truth.ω[1][1]
 
-initialize!(mech, :box, x=x2, v=v15, q=q2, ω=ϕ15)
+initialize!(mech, :block, x=x2, v=v15, q=q2, ω=ϕ15)
 traj_sim = simulate!(mech, 0.80, record=true,
     opts=SolverOptions(btol=1e-6, rtol=1e-6, verbose=false))
 
diff --git a/examples/system_identification/learning.jl b/examples/system_identification/learning.jl
@@ -22,7 +22,7 @@ open(vis)
 S = 7
 timestep= 1/148 * S
 gscaled = -9.81*20
-mech = get_mechanism(:box, 
+mech = get_mechanism(:block, 
 	timestep=timestep, 
 	gravity=gravityscaled, 
 	friction_coefficient=0.2, 
@@ -31,7 +31,7 @@ mech = get_mechanism(:box,
 	mode=:box);
 
 # ## Simualate
-initialize!(mech, :box, 
+initialize!(mech, :block, 
 	x=[0.0, -1.0, 1.0], 
 	v=[0.0, 2.0, 1.0], 
 	ω=[2.0, 5.0, 10.0])
@@ -61,18 +61,18 @@ data0 = params0[:data]
 ################################################################################
 ## Optimization Objective: Evaluation & Gradient
 ################################################################################
-clean_loss(:box, pairs0, data0, 
+clean_loss(:block, pairs0, data0, 
 	n_sample=250,
 	opts=SolverOptions(btol=3e-4, rtol=3e-4), 
 	timestep=timestep, 
 	gravity=gravityscaled)
 
-[clean_loss(:box, pairs0, data0 + [i;zeros(55)],
+[clean_loss(:block, pairs0, data0 + [i;zeros(55)],
 	opts=SolverOptions(btol=3e-4, rtol=3e-4), 
 	timestep=timestep, 
 	gravity=gravityscaled)
 	for i in Vector(-0.10:0.01:0.1)]
-[clean_loss(:box, pairs0, data0 + [0;i;zeros(54)],
+[clean_loss(:block, pairs0, data0 + [0;i;zeros(54)],
 	opts=SolverOptions(btol=3e-4, rtol=3e-4), 
 	timestep=timestep, 
 	gravity=gravityscaled)
@@ -108,7 +108,7 @@ lower = [0.00, 0.05]
 upper = [0.80, 1.50]
 
 function f0(d; rot=0)
-	return clean_loss(:box, pairs0, d2data(d), 
+	return clean_loss(:block, pairs0, d2data(d), 
 		n_sample=200, 
 		timestep=timestep, 
 		gravity=gravityscaled, 
@@ -118,7 +118,7 @@ end
 
 function fgH0(d; 
 	rot=0)
-	f, g, H = clean_loss(:box, pairs0, d2data(d), 
+	f, g, H = clean_loss(:block, pairs0, d2data(d), 
 		n_sample=200, 
 		timestep=timestep, 
 		gravity=gravityscaled, 
@@ -151,7 +151,7 @@ dsol #[0.171, 0.933]
 using Plots; pyplot()
 x=range(0.00,stop=0.80,length=10)
 y=range(0.10,stop=1.50,length=10)
-f(x,y) = log(10, clean_loss(:box, pairs0, d2data([x,y]), n_sample=50, timestep=timestep, gravity=gravityscaled,
+f(x,y) = log(10, clean_loss(:block, pairs0, d2data([x,y]), n_sample=50, timestep=timestep, gravity=gravityscaled,
 	opts=SolverOptions(btol=3e-4, rtol=3e-4))[1])
 plot(x,y,f,st=:surface,camera=(25,75))
 
@@ -212,11 +212,11 @@ upper = [0.80,
 	-0.05, -0.05, +1.50]
 
 function f0(d; rot=0, n_sample=50)
-	return clean_loss(:box, pairs0, d2data(d), n_sample=n_sample, timestep=timestep, gravity=gravityscaled, rot=rot, opts=SolverOptions(btol=3e-4, rtol=3e-4))[1]
+	return clean_loss(:block, pairs0, d2data(d), n_sample=n_sample, timestep=timestep, gravity=gravityscaled, rot=rot, opts=SolverOptions(btol=3e-4, rtol=3e-4))[1]
 end
 
 function fgH0(d; rot=0, n_sample=50)
-	f, g, H = clean_loss(:box, pairs0, d2data(d), n_sample=n_sample, timestep=timestep, gravity=gravityscaled, rot=rot, opts=SolverOptions(btol=3e-4, rtol=3e-4))
+	f, g, H = clean_loss(:block, pairs0, d2data(d), n_sample=n_sample, timestep=timestep, gravity=gravityscaled, rot=rot, opts=SolverOptions(btol=3e-4, rtol=3e-4))
 	return f, ∇d2data' * g, ∇d2data' * H * ∇d2data
 end
 dsol, Dsol = quasi_newton_solve(f0, fgH0, d0, iter=50, gtol=1e-8, ftol=1e-6,
diff --git a/examples/system_identification/synthetic_examples/box.jl b/examples/system_identification/synthetic_examples/box.jl
@@ -11,8 +11,8 @@ open(vis)
 # Include new files
 include( "../utils.jl")
 
-mech = get_mechanism(:box, timestep=0.05, g=-9.81, friction_coefficient=0.2, radius=0.00, side=0.50, mode=:box);
-initialize!(mech, :box, x=[0,-1,1.], v=[0,2,1.], ω=[2,5,10.])
+mech = get_mechanism(:block, timestep=0.05, g=-9.81, friction_coefficient=0.2, radius=0.00, side=0.50, mode=:box);
+initialize!(mech, :block, x=[0,-1,1.], v=[0,2,1.], ω=[2,5,10.])
 storage = simulate!(mech, 5.0, record=true,
     opts=SolverOptions(btol=1e-6, rtol=1e-6, verbose=false))
 visualize(mech, storage, vis=vis, show_contact=false)
@@ -25,7 +25,7 @@ init_kwargs = Dict(:xlims => [[0,0,0.2], [1,1,0.4]],
 				   :vlims => [-2ones(3), [2,2,-1.]],
 				   :ωlims => [-6ones(3), 6ones(3)])
 mech_kwargs = Dict(:friction_coefficient => 0.1, :radius => 0.0, :side => 0.5)
-generate_dataset(:box, H=0.40, N=35,
+generate_dataset(:block, H=0.40, N=35,
 	opts=SolverOptions(btol=3e-4, rtol=3e-4),
 	init_kwargs=init_kwargs,
 	mech_kwargs=mech_kwargs,
@@ -36,20 +36,20 @@ generate_dataset(:box, H=0.40, N=35,
 ################################################################################
 # Load Dataset
 ################################################################################
-params0, trajs0, pairs0 = open_dataset(:box; N = 35, mech_kwargs...)
+params0, trajs0, pairs0 = open_dataset(:block; N = 35, mech_kwargs...)
 
 data0 = params0[:data]
 
 ################################################################################
 # Optimization Objective: Evaluation & Gradient
 ################################################################################
-# @benchmark clean_loss(:box, pairs0, data0, opts=SolverOptions(btol=3e-4, rtol=3e-4))
-# @profiler clean_loss(:box, pairs0, data0, opts=SolverOptions(btol=3e-4, rtol=3e-4))
-clean_loss(:box, pairs0, data0, opts=SolverOptions(btol=3e-4, rtol=3e-4))
+# @benchmark clean_loss(:block, pairs0, data0, opts=SolverOptions(btol=3e-4, rtol=3e-4))
+# @profiler clean_loss(:block, pairs0, data0, opts=SolverOptions(btol=3e-4, rtol=3e-4))
+clean_loss(:block, pairs0, data0, opts=SolverOptions(btol=3e-4, rtol=3e-4))
 
-[clean_loss(:box, pairs0, data0 + [i;zeros(55)], opts=SolverOptions(btol=3e-4, rtol=3e-4))
+[clean_loss(:block, pairs0, data0 + [i;zeros(55)], opts=SolverOptions(btol=3e-4, rtol=3e-4))
 	for i in Vector(-0.10:0.01:0.1)]
-[clean_loss(:box, pairs0, data0 + [0;i;zeros(54)], opts=SolverOptions(btol=3e-4, rtol=3e-4))
+[clean_loss(:block, pairs0, data0 + [0;i;zeros(54)], opts=SolverOptions(btol=3e-4, rtol=3e-4))
 	for i in Vector(-0.10:0.01:0.1)]
 
 plot(hcat([p[1][1:3] for p in pairs0]...)')
@@ -109,11 +109,11 @@ upper = [0.80,
 	-0.05, -0.05, +1.00]
 
 function f0(d; rot=0)
-	return clean_loss(:box, pairs0, d2data(d), n_sample=35, rot=rot, opts=SolverOptions(btol=3e-4, rtol=3e-4))[1]
+	return clean_loss(:block, pairs0, d2data(d), n_sample=35, rot=rot, opts=SolverOptions(btol=3e-4, rtol=3e-4))[1]
 end
 
 function fgH0(d; rot=0)
-	f, g, H = clean_loss(:box, pairs0, d2data(d), n_sample=35, rot=rot, opts=SolverOptions(btol=3e-4, rtol=3e-4))
+	f, g, H = clean_loss(:block, pairs0, d2data(d), n_sample=35, rot=rot, opts=SolverOptions(btol=3e-4, rtol=3e-4))
 	return f, ∇d2data' * g, ∇d2data' * H * ∇d2data
 end
 dsol = quasi_newton_solve(f0, fgH0, d0, iter=50, gtol=1e-8, ftol=1e-6,
@@ -178,11 +178,11 @@ upper = [
 	0.80, -0.05, -0.05, +1.00]
 
 function f0(d; rot=0)
-	return clean_loss(:box, pairs0, d2data(d), n_sample=35, rot=rot, opts=SolverOptions(btol=3e-4, rtol=3e-4))[1]
+	return clean_loss(:block, pairs0, d2data(d), n_sample=35, rot=rot, opts=SolverOptions(btol=3e-4, rtol=3e-4))[1]
 end
 
 function fgH0(d; rot=0)
-	f, g, H = clean_loss(:box, pairs0, d2data(d), n_sample=35, rot=rot, opts=SolverOptions(btol=3e-4, rtol=3e-4))
+	f, g, H = clean_loss(:block, pairs0, d2data(d), n_sample=35, rot=rot, opts=SolverOptions(btol=3e-4, rtol=3e-4))
 	return f, ∇d2data' * g, ∇d2data' * H * ∇d2data
 end
 
diff --git a/test/behaviors.jl b/test/behaviors.jl
@@ -18,12 +18,12 @@ end
 
 @testset "Box toss" begin
     for timestep in [0.10, 0.05, 0.01, 0.005]
-        mech = get_mechanism(:box,
+        mech = get_mechanism(:block,
             timestep=timestep,
             gravity=-9.81,
             friction_coefficient = 0.1)
 
-        initialize!(mech, :box,
+        initialize!(mech, :block,
             x=[0.0, 0.0, 0.5],
             v=[1.0, 1.5, 1.0],
             ω=[5.0, 4.0, 2.0] .* timestep)
diff --git a/test/data.jl b/test/data.jl
@@ -141,7 +141,7 @@ end
 	for (spring, damper) in [(0.0, 0.0), (2.0, 0.3)]
 		test_data_system(:sphere, 
 			contact=false)
-		test_data_system(:box, 
+		test_data_system(:block, 
 			contact=false)
 		test_data_system(:box2D, 
 			contact=false)
@@ -212,7 +212,7 @@ end
 	for (spring, damper) in [(0.0, 0.0), (2.0, 0.3)]
 		test_data_system(:sphere, 
 			contact=true)
-		test_data_system(:box, 
+		test_data_system(:block, 
 			contact=true)
 		test_data_system(:box2D, 
 			contact=true)
diff --git a/test/energy.jl b/test/energy.jl
@@ -98,14 +98,14 @@ end
 @testset "Dice" begin
     gravity0 = -10.0
 
-    mech = get_mechanism(:box, 
+    mech = get_mechanism(:block, 
         timestep=timestep0, 
         gravity=gravity0, 
         contact=false)
 
     v0 = [1,2,3.0]
     ω0 = [1,1,1.0]
-    initialize!(mech, :box, 
+    initialize!(mech, :block, 
         v=v0, 
         ω=ω0)
 
diff --git a/test/environments.jl b/test/environments.jl
diff --git a/test/jacobian.jl b/test/jacobian.jl
diff --git a/test/mechanisms.jl b/test/mechanisms.jl
diff --git a/test/minimal.jl b/test/minimal.jl
diff --git a/test/momentum.jl b/test/momentum.jl
