Skip to content

Commit

Permalink
Upgrade rotations (#14)
Browse files Browse the repository at this point in the history 
* Upgrade Rotations

* Update formatting

---------

Co-authored-by: Brian Jackson <[email protected]>
  • Loading branch information
@bjack205
bjack205 and Brian Jackson authored Mar 20, 2024
1 parent ed2b469 commit 8a2bb86
Show file tree
Hide file tree
Showing 11 changed files with 377 additions and 373 deletions.
4 changes: 2 additions & 2 deletions Project.toml
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "RobotZoo"
uuid = "74be38bb-dcc2-4b9e-baf3-d6373cd95f10"
authors = ["Brian Jackson <[email protected]>"]
version = "0.3.2"
version = "0.3.3"

[deps]
FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
Expand All @@ -14,7 +14,7 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
[compat]
FiniteDiff = "2.10"
ForwardDiff = "0.10"
RobotDynamics = "0.4.1"
RobotDynamics = "0.4.8"
Rotations = "1.0"
StaticArrays = "0.12, 1.0"
julia = "1"
12 changes: 6 additions & 6 deletions src/cartpole.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ with keyword arguments
* `l` - length of the pendulum, in m (default = 0.5)
* `g` - gravity, in m/s² (default = 9.81)
"""
@autodiff struct Cartpole{T} <: ContinuousDynamics
@autodiff struct Cartpole{T} <: ContinuousDynamics
mc::T
mp::T
l::T
Expand All @@ -32,18 +32,18 @@ function dynamics(model::Cartpole, x, u)
l = model.l # length of the pole in m
g = model.g # gravity m/s^2

q = x[ @SVector [1,2] ]
qd = x[ @SVector [3,4] ]
q = x[@SVector [1, 2]]
qd = x[@SVector [3, 4]]

s = sin(q[2])
c = cos(q[2])

H = @SMatrix [mc+mp mp*l*c; mp*l*c mp*l^2]
C = @SMatrix [0 -mp*qd[2]*l*s; 0 0]
G = @SVector [0, mp*g*l*s]
G = @SVector [0, mp * g * l * s]
B = @SVector [1, 0]

qdd = -H\(C*qd + G - B*u[1])
qdd = -H \ (C * qd + G - B * u[1])
return [qd; qdd]
end

Expand All @@ -55,4 +55,4 @@ RobotDynamics.state_dim(::Cartpole) = 4
RobotDynamics.control_dim(::Cartpole) = 1

# Base.position(::Cartpole, x::StaticVector) = SA[0,x[1],0]
# RobotDynamics.orientation(::Cartpole, x::StaticVector) = UnitQuaternion(expm(SA[1,0,0]*x[2]))
# RobotDynamics.orientation(::Cartpole, x::StaticVector) = QuatRotation(expm(SA[1,0,0]*x[2]))
124 changes: 62 additions & 62 deletions src/copymodel.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,20 +7,20 @@ struct CopyModel{K,N,M,L} <: ContinuousDynamics
uinds::Vector{SVector{M,Int}}
end

function CopyModel(model::L, K::Int) where L <: ContinuousDynamics
n,m = size(model)
function CopyModel(model::L, K::Int) where {L<:ContinuousDynamics}
n, m = size(model)
xind = SVector{n}(1:n)
uind = SVector{m}(1:m)
xinds = [xind .+ (i-1)*n for i = 1:K]
uinds = [uind .+ (i-1)*m for i = 1:K]
xinds = [xind .+ (i - 1) * n for i = 1:K]
uinds = [uind .+ (i - 1) * m for i = 1:K]
ix = SVector{n}(1:n)
iu = SVector{m}(n .+ (1:m))
return CopyModel{K,n,m,L}(model, ix, iu, xinds, uinds)
end

Base.size(::CopyModel{K,N,M}) where {K,N,M} = N*K,M*K
Base.size(::CopyModel{K,N,M}) where {K,N,M} = N * K, M * K

function build_state(model::CopyModel{K}, xs...) where K
function build_state(model::CopyModel{K}, xs...) where {K}
@assert length(xs) == K
x = xs[1]
for i = 2:K
Expand All @@ -30,23 +30,23 @@ function build_state(model::CopyModel{K}, xs...) where K
end

function Base.zeros(model::CopyModel{K,N,M}) where {K,N,M}
x,u = zeros(model.model)
return repeat(x,K), repeat(u,K)
x, u = zeros(model.model)
return repeat(x, K), repeat(u, K)
end

function Base.rand(model::CopyModel{K}) where K
x,u = rand(model.model)
function Base.rand(model::CopyModel{K}) where {K}
x, u = rand(model.model)
for i = 2:K
x_,u_ = rand(model.model)
x_, u_ = rand(model.model)
x = [x; x_]
u = [u; u_]
end
return x,u
return x, u
end


function states(model::CopyModel{K,N,M}, x) where {K,N,M}
reshape(x,N,K)
reshape(x, N, K)
return [x[i] for i in model.xinds]
end

Expand All @@ -61,7 +61,7 @@ function states(model::CopyModel, Z::Traj, k::Int)
end

function controls(model::CopyModel{K,N,M}, u) where {K,N,M}
reshape(u,M,K)
reshape(u, M, K)
end

controls(model::CopyModel{K,L}, z::KnotPoint, k::Int) where {L,K} =
Expand All @@ -80,11 +80,11 @@ function TrajectoryOptimization.get_trajectory(model::CopyModel, Z::Traj, k::Int
map(Z) do z
x = state(z)[xinds[k]]
u = control(z)[uinds[k]]
KnotPoint(x,u,z.dt,z.t)
KnotPoint(x, u, z.dt, z.t)
end
end

function TrajectoryOptimization.dynamics(model::CopyModel{K}, x, u, t=0.0) where K
function TrajectoryOptimization.dynamics(model::CopyModel{K}, x, u, t=0.0) where {K}
xinds = model.xinds
uinds = model.uinds
xdot = dynamics(model.model, x[xinds[1]], u[uinds[1]], t)
Expand All @@ -96,10 +96,10 @@ function TrajectoryOptimization.dynamics(model::CopyModel{K}, x, u, t=0.0) where
end

function TrajectoryOptimization.discrete_dynamics(::Type{Q}, model::CopyModel{K},
z::KnotPoint) where {K,Q <: TrajectoryOptimization.Implicit}
z::KnotPoint) where {K,Q<:TrajectoryOptimization.Implicit}
xinds = model.xinds
uinds = model.uinds
x,u = state(z), control(z)
x, u = state(z), control(z)
xdot = discrete_dynamics(Q, model.model, x[xinds[1]], u[uinds[1]], z.t, z.dt)
for i = 2:K
xdot_ = discrete_dynamics(Q, model.model, x[xinds[i]], u[uinds[i]], z.t, z.dt)
Expand All @@ -109,33 +109,33 @@ function TrajectoryOptimization.discrete_dynamics(::Type{Q}, model::CopyModel{K}
end

function TrajectoryOptimization.jacobian(model::CopyModel{K,N0,M0},
z::KnotPoint{T,N,M,NM}) where {K,N0,M0,T,N,M,NM}
A0 = @SMatrix zeros(N0,N0)
B0 = @SMatrix zeros(N0,M0)
z::KnotPoint{T,N,M,NM}) where {K,N0,M0,T,N,M,NM}
A0 = @SMatrix zeros(N0, N0)
B0 = @SMatrix zeros(N0, M0)
ix = model.ix
iu = model.iu

xinds = model.xinds
uinds = model.uinds
x,u = state(z), control(z)
x, u = state(z), control(z)

# Process fist model
z_ = [x[xinds[1]]; u[uinds[1]]]
z_ = TrajectoryOptimization.StaticKnotPoint(z_, ix, iu, z.dt, z.t)
∇f = jacobian(model.model, z_)
A,B = ∇f[ix,ix], ∇f[ix,iu]
A, B = ∇f[ix, ix], ∇f[ix, iu]
# append zeros after
for i = 2:K
A = [A A0]
B = [B B0]
A = [A A0]
B = [B B0]
end

# loop over the rest of the models, appending below the first
for i = 2:K
z_ = [x[xinds[i]]; u[uinds[i]]]
z_ = TrajectoryOptimization.StaticKnotPoint(z_, ix, iu, z.dt, z.t)
∇f = jacobian(model.model, z_)
A_,B_ = ∇f[ix,ix], ∇f[ix,iu]
A_, B_ = ∇f[ix, ix], ∇f[ix, iu]

# prepend zeros
for j = 1:i-1
Expand Down Expand Up @@ -195,33 +195,33 @@ end
# end

function discrete_jacobian!(::Type{Q}, ∇f, model::Dynamics.CopyModel{K,N0,M0},
z::KnotPoint{T,N,M,NM}) where {T,N,M,NM,K,N0,M0,Q<:TrajectoryOptimization.Implicit}
z::KnotPoint{T,N,M,NM}) where {T,N,M,NM,K,N0,M0,Q<:TrajectoryOptimization.Implicit}
xinds = model.xinds
uinds = model.uinds
ix = model.ix
iu = model.iu
iz = [ix; iu; @SVector [N0+M0+1]]
x,u = state(z), control(z)
for i = 1:K
z_ = [x[xinds[i]]; u[uinds[i]]]
z_ = StaticKnotPoint(z_, ix, iu, z.dt, z.t)
∇f_ = discrete_jacobian(Q, model.model, z_)
A,B,C = ∇f_[ix,ix], ∇f_[ix,iu], ∇f_[ix,N0+M0+1]
ix = model.ix
iu = model.iu
iz = [ix; iu; @SVector [N0 + M0 + 1]]
x, u = state(z), control(z)
for i = 1:K
z_ = [x[xinds[i]]; u[uinds[i]]]
z_ = StaticKnotPoint(z_, ix, iu, z.dt, z.t)
∇f_ = discrete_jacobian(Q, model.model, z_)
A, B, C = ∇f_[ix, ix], ∇f_[ix, iu], ∇f_[ix, N0+M0+1]

ix_ = ix .+ (i-1)*N0
iu_ = iu .+ (i-1)*M0 .+ N .- N0
it_ = NM+1
∇f[ix_,ix_] .= A
∇f[ix_,iu_] .= B
∇f[ix_,it_] .= C
end
ix_ = ix .+ (i - 1) * N0
iu_ = iu .+ (i - 1) * M0 .+ N .- N0
it_ = NM + 1
∇f[ix_, ix_] .= A
∇f[ix_, iu_] .= B
∇f[ix_, it_] .= C
end
end

function state_diff_size(model::CopyModel{K}) where K
return state_diff_size(model.model)*K
function state_diff_size(model::CopyModel{K}) where {K}
return state_diff_size(model.model) * K
end

state_diff(model::CopyModel, x::SVector, x0::SVector) = x-x0
state_diff(model::CopyModel, x::SVector, x0::SVector) = x - x0
function state_diff(model::CopyModel{K,N,M,L}, x::SVector, x0::SVector) where {K,N,M,L<:RigidBody}
xinds = model.xinds
dx = state_diff(model.model, x[xinds[1]], x0[xinds[1]])
Expand All @@ -234,7 +234,7 @@ end

TrajectoryOptimization.state_diff_jacobian!(G, model::CopyModel, Z::Traj) = nothing
function TrajectoryOptimization.state_diff_jacobian!(G,
model::CopyModel{<:Any,<:Any,<:Any,L}, Z::Traj) where {L<:RigidBody{R}} where R
model::CopyModel{<:Any,<:Any,<:Any,L}, Z::Traj) where {L<:RigidBody{R}} where {R}
for k in eachindex(Z)
G[k] .= state_diff_jacobian(model, state(Z[k]))
end
Expand All @@ -243,15 +243,15 @@ end

state_diff_jacobian(::CopyModel, x::SVector) = I
@generated function state_diff_jacobian(model::CopyModel{K,N,M,L},
x::SVector) where {K,N,M,L<:RigidBody{R}} where R
if R <: UnitQuaternion
if R <: UnitQuaternion{T,IdentityMap} where T
x::SVector) where {K,N,M,L<:RigidBody{R}} where {R}
if R <: QuatRotation
if R <: QuatRotation{T,IdentityMap} where {T}
return :(I)
else
G0 = @SMatrix zeros(N,N-1)
G0 = @SMatrix zeros(N, N - 1)
end
else
G0 = @SMatrix zeros(N,N)
G0 = @SMatrix zeros(N, N)
end
quote
xinds = model.xinds
Expand All @@ -274,27 +274,27 @@ state_diff_jacobian(::CopyModel, x::SVector) = I
end

@generated function TrajectoryOptimization.∇²differential(model::CopyModel{K,N,M,L},
x::SVector, dx::SVector) where {K,N,M,L<:RigidBody{R}} where R
x::SVector, dx::SVector) where {K,N,M,L<:RigidBody{R}} where {R}

if R <: UnitQuaternion
if R <: UnitQuaternion{T,IdentityMap} where T
return :(I*0)
if R <: QuatRotation
if R <: QuatRotation{T,IdentityMap} where {T}
return :(I * 0)
else
G0 = @SMatrix zeros(N-1,N-1)
G0 = @SMatrix zeros(N - 1, N - 1)
end
else
G0 = @SMatrix zeros(N,N)
G0 = @SMatrix zeros(N, N)
end
quote
dx_ = reshape(dx,:,K)
dx_ = reshape(dx, :, K)
ix_ = SVector{12}(1:12)
xinds = model.xinds
G = TrajectoryOptimization.∇²differential(model.model, x[xinds[1]], dx[ix_])
for i = 2:K
G = [G $G0]
end
for i = 2:K
dx_ = dx[ix_ .+ (i-1)*12]
dx_ = dx[ix_.+(i-1)*12]
G_ = TrajectoryOptimization.∇²differential(model.model, x[xinds[i]], dx_)
for j = 1:i-1
G_ = [$G0 G_]
Expand All @@ -309,11 +309,11 @@ end
end

function ∇²differential!(G, model::Dynamics.CopyModel{K,N,M,L},
x::SVector, dx::Vector) where {K,N,M,L<:RigidBody{R}} where R
x::SVector, dx::Vector) where {K,N,M,L<:RigidBody{R}} where {R}
ix = SVector{12}(1:12)
xinds = model.xinds
for i = 1:K
ix_ = ix .+ (i-1)*12
ix_ = ix .+ (i - 1) * 12
dx_ = dx[ix_]
G_ = ∇²differential(model.model, x[xinds[i]], dx_)
G[ix_, ix_] .= G_
Expand Down
8 changes: 4 additions & 4 deletions src/double_integrator.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ end
function DoubleIntegrator(D=1; gravity=zeros(D))
pos = SVector{D,Int}(1:D)
vel = SVector{D,Int}(D .+ (1:D))
DoubleIntegrator{2D,D}(pos,vel, gravity)
DoubleIntegrator{2D,D}(pos, vel, gravity)
end

RobotDynamics.state_dim(::DoubleIntegrator{N,M}) where {N,M} = N
Expand All @@ -29,11 +29,11 @@ RobotDynamics.control_dim(::DoubleIntegrator{N,M}) where {N,M} = M
@generated function dynamics(di::DoubleIntegrator{N,M}, x, u) where {N,M}
vel = [:(x[$i]) for i = M+1:N]
us = [:(u[$i] + di.gravity[$i]) for i = 1:M]
:(SVector{$N}($(vel...),$(us...)))
:(SVector{$N}($(vel...), $(us...)))
end
@generated function dynamics!(di::DoubleIntegrator{N,M}, xdot, x, u) where {N,M}
vel = [:(xdot[$i-M] = x[$i]) for i = M+1:N]
us = [:(xdot[$(i+M)] = u[$i] + di.gravity[$i]) for i = 1:M]
us = [:(xdot[$(i + M)] = u[$i] + di.gravity[$i]) for i = 1:M]
quote
$(vel...)
$(us...)
Expand All @@ -42,4 +42,4 @@ end
end

Base.position(::DoubleIntegrator{<:Any,2}, x) = @SVector [x[1], x[2], 0]
orientation(::DoubleIntegrator, x) = UnitQuaternion(I)
orientation(::DoubleIntegrator, x) = QuatRotation(I)
Loading

2 comments on commit 8a2bb86

@bjack205
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator register

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request created: JuliaRegistries/General/103299

Tip: Release Notes

Did you know you can add release notes too? Just add markdown formatted text underneath the comment after the text
"Release notes:" and it will be added to the registry PR, and if TagBot is installed it will also be added to the
release that TagBot creates. i.e.

@JuliaRegistrator register

Release notes:

## Breaking changes

- blah

To add them here just re-invoke and the PR will be updated.

Tagging

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a v0.3.3 -m "<description of version>" 8a2bb862a64cbfa00c0ff5e391e38db09d89ab8d
git push origin v0.3.3

Also, note the warning: Version 0.3.3 skips over 0.3.2
This can be safely ignored. However, if you want to fix this you can do so. Call register() again after making the fix. This will update the Pull request.

Please sign in to comment.